diff --git a/pkg/connector/auth_recovery.go b/pkg/connector/auth_recovery.go index 9277ea5..c818614 100644 --- a/pkg/connector/auth_recovery.go +++ b/pkg/connector/auth_recovery.go @@ -14,6 +14,10 @@ type lineCallDeps[T any] struct { call func(*line.Client) (T, error) } +var recoverLineToken = func(lc *LineClient, ctx context.Context) error { + return lc.recoverToken(ctx) +} + func callLineWithRecovery[T any](ctx context.Context, client *line.Client, deps lineCallDeps[T]) (*line.Client, T, error) { if client == nil { client = deps.newClient() @@ -53,7 +57,7 @@ func (lc *LineClient) callLine(ctx context.Context, call func(*line.Client) erro func (lc *LineClient) callLineUsing(ctx context.Context, client *line.Client, call func(*line.Client) error) (*line.Client, error) { client, _, err := callLineWithRecovery(ctx, client, lineCallDeps[struct{}]{ newClient: func() *line.Client { return lc.newClient() }, - recover: lc.recoverToken, + recover: func(ctx context.Context) error { return recoverLineToken(lc, ctx) }, isAuthError: lc.isTokenError, call: func(client *line.Client) (struct{}, error) { return struct{}{}, call(client) @@ -72,7 +76,7 @@ func callLineResult[T any](lc *LineClient, ctx context.Context, call func(*line. func callLineResultUsing[T any](lc *LineClient, ctx context.Context, client *line.Client, call func(*line.Client) (T, error)) (*line.Client, T, error) { client, res, err := callLineWithRecovery(ctx, client, lineCallDeps[T]{ newClient: func() *line.Client { return lc.newClient() }, - recover: lc.recoverToken, + recover: func(ctx context.Context) error { return recoverLineToken(lc, ctx) }, isAuthError: lc.isTokenError, call: call, }) diff --git a/pkg/connector/e2ee_keys.go b/pkg/connector/e2ee_keys.go index 65c41c0..6326f6b 100644 --- a/pkg/connector/e2ee_keys.go +++ b/pkg/connector/e2ee_keys.go @@ -13,6 +13,15 @@ import ( const noE2EETTL = 1 * time.Hour +var ( + negotiateE2EEPublicKeyWithClient = func(client *line.Client, mid string) (*line.E2EEPublicKey, error) { + return client.NegotiateE2EEPublicKey(mid) + } + getE2EEPublicKeyWithClient = func(client *line.Client, mid string, keyVersion, keyID int) (*line.E2EEPublicKey, error) { + return client.GetE2EEPublicKey(mid, keyVersion, keyID) + } +) + // fetchAndUnwrapGroupKey retrieves a specific group key (or the latest when groupKeyID == 0) // and unwraps it so the E2EE manager can encrypt/decrypt group messages. // If no group key exists yet (TalkException code 5), it auto-registers one and retries. @@ -86,7 +95,7 @@ func (lc *LineClient) fetchAndUnwrapGroupKey(ctx context.Context, chatMid string return nil } -func (lc *LineClient) ensurePeerKey(_ context.Context, mid string) (int, string, error) { +func (lc *LineClient) ensurePeerKey(ctx context.Context, mid string) (int, string, error) { lc.cacheMu.Lock() if lc.peerKeys == nil { lc.peerKeys = make(map[string]peerKeyInfo) @@ -107,8 +116,9 @@ func (lc *LineClient) ensurePeerKey(_ context.Context, mid string) (int, string, return cached.raw, cached.pub, nil } } - client := lc.newClient() - res, err := client.NegotiateE2EEPublicKey(mid) + _, res, err := callLineResult(lc, ctx, func(client *line.Client) (*line.E2EEPublicKey, error) { + return negotiateE2EEPublicKeyWithClient(client, mid) + }) if err != nil { // Cache negative result so we don't keep hitting the API if line.IsNoUsableE2EEPublicKey(err) { @@ -305,7 +315,7 @@ func (lc *LineClient) getGroupMemberMIDsViaMatrix(ctx context.Context, chatMid s return mids, nil } -func (lc *LineClient) ensurePeerKeyByID(_ context.Context, mid string, keyID int) (int, string, error) { +func (lc *LineClient) ensurePeerKeyByID(ctx context.Context, mid string, keyID int) (int, string, error) { lc.cacheMu.Lock() if lc.peerKeys == nil { lc.peerKeys = make(map[string]peerKeyInfo) @@ -319,9 +329,10 @@ func (lc *LineClient) ensurePeerKeyByID(_ context.Context, mid string, keyID int return cached.raw, cached.pub, nil } - client := lc.newClient() // keyVersion 1 - res, err := client.GetE2EEPublicKey(mid, 1, keyID) + _, res, err := callLineResult(lc, ctx, func(client *line.Client) (*line.E2EEPublicKey, error) { + return getE2EEPublicKeyWithClient(client, mid, 1, keyID) + }) if err != nil { return 0, "", err } diff --git a/pkg/connector/e2ee_keys_test.go b/pkg/connector/e2ee_keys_test.go new file mode 100644 index 0000000..cb9d8a5 --- /dev/null +++ b/pkg/connector/e2ee_keys_test.go @@ -0,0 +1,162 @@ +package connector + +import ( + "context" + "encoding/json" + "errors" + "io" + "testing" + + "github.com/rs/zerolog" + "maunium.net/go/mautrix/bridgev2" + + "github.com/highesttt/matrix-line-messenger/pkg/line" +) + +func newPeerKeyTestClient() *LineClient { + return &LineClient{ + AccessToken: "access", + UserLogin: &bridgev2.UserLogin{ + Bridge: &bridgev2.Bridge{Log: zerolog.New(io.Discard)}, + }, + } +} + +func TestEnsurePeerKeyRecoversAndRetriesRefreshRequired(t *testing.T) { + oldNewClient := newLineAPIClient + oldNegotiate := negotiateE2EEPublicKeyWithClient + oldRecover := recoverLineToken + t.Cleanup(func() { + newLineAPIClient = oldNewClient + negotiateE2EEPublicKeyWithClient = oldNegotiate + recoverLineToken = oldRecover + }) + + var newClientCalls int + newLineAPIClient = func(token string) *line.Client { + newClientCalls++ + return line.NewClient(token) + } + + var negotiateCalls int + negotiateE2EEPublicKeyWithClient = func(*line.Client, string) (*line.E2EEPublicKey, error) { + negotiateCalls++ + if negotiateCalls == 1 { + return nil, errAuthRequired + } + return &line.E2EEPublicKey{ + KeyID: json.Number("42"), + PublicKey: "peer-public-key", + }, nil + } + + var recoverCalls int + recoverLineToken = func(*LineClient, context.Context) error { + recoverCalls++ + return nil + } + + lc := newPeerKeyTestClient() + keyID, publicKey, err := lc.ensurePeerKey(context.Background(), "peer-mid") + if err != nil { + t.Fatalf("ensurePeerKey returned error: %v", err) + } + if keyID != 42 || publicKey != "peer-public-key" { + t.Fatalf("ensurePeerKey returned keyID=%d publicKey=%q", keyID, publicKey) + } + if negotiateCalls != 2 { + t.Fatalf("negotiate calls = %d, want 2", negotiateCalls) + } + if recoverCalls != 1 { + t.Fatalf("recovery calls = %d, want 1", recoverCalls) + } + if newClientCalls != 2 { + t.Fatalf("new clients = %d, want 2", newClientCalls) + } +} + +func TestEnsurePeerKeyByIDRecoversAndRetriesRefreshRequired(t *testing.T) { + oldNewClient := newLineAPIClient + oldGetKey := getE2EEPublicKeyWithClient + oldRecover := recoverLineToken + t.Cleanup(func() { + newLineAPIClient = oldNewClient + getE2EEPublicKeyWithClient = oldGetKey + recoverLineToken = oldRecover + }) + + var newClientCalls int + newLineAPIClient = func(token string) *line.Client { + newClientCalls++ + return line.NewClient(token) + } + + var getKeyCalls int + getE2EEPublicKeyWithClient = func(*line.Client, string, int, int) (*line.E2EEPublicKey, error) { + getKeyCalls++ + if getKeyCalls == 1 { + return nil, errAuthRequired + } + return &line.E2EEPublicKey{ + KeyID: json.Number("5910969"), + PublicKey: "specific-peer-public-key", + }, nil + } + + var recoverCalls int + recoverLineToken = func(*LineClient, context.Context) error { + recoverCalls++ + return nil + } + + lc := newPeerKeyTestClient() + keyID, publicKey, err := lc.ensurePeerKeyByID(context.Background(), "peer-mid", 5910969) + if err != nil { + t.Fatalf("ensurePeerKeyByID returned error: %v", err) + } + if keyID != 5910969 || publicKey != "specific-peer-public-key" { + t.Fatalf("ensurePeerKeyByID returned keyID=%d publicKey=%q", keyID, publicKey) + } + if getKeyCalls != 2 { + t.Fatalf("get key calls = %d, want 2", getKeyCalls) + } + if recoverCalls != 1 { + t.Fatalf("recovery calls = %d, want 1", recoverCalls) + } + if newClientCalls != 2 { + t.Fatalf("new clients = %d, want 2", newClientCalls) + } +} + +func TestEnsurePeerKeyCachesNoUsablePublicKeyWithoutRecovery(t *testing.T) { + oldNegotiate := negotiateE2EEPublicKeyWithClient + oldRecover := recoverLineToken + t.Cleanup(func() { + negotiateE2EEPublicKeyWithClient = oldNegotiate + recoverLineToken = oldRecover + }) + + var negotiateCalls int + negotiateE2EEPublicKeyWithClient = func(*line.Client, string) (*line.E2EEPublicKey, error) { + negotiateCalls++ + return nil, line.ErrNoUsableE2EEPublicKey + } + recoverLineToken = func(*LineClient, context.Context) error { + t.Fatal("recovery should not be called for no-usable-public-key errors") + return nil + } + + lc := newPeerKeyTestClient() + _, _, err := lc.ensurePeerKey(context.Background(), "peer-mid") + if !errors.Is(err, line.ErrNoUsableE2EEPublicKey) { + t.Fatalf("ensurePeerKey error = %v, want ErrNoUsableE2EEPublicKey", err) + } + + _, _, err = lc.ensurePeerKey(context.Background(), "peer-mid") + if !errors.Is(err, line.ErrNoUsableE2EEPublicKey) { + t.Fatalf("cached ensurePeerKey error = %v, want ErrNoUsableE2EEPublicKey", err) + } + if negotiateCalls != 1 { + t.Fatalf("negotiate calls = %d, want cached negative lookup", negotiateCalls) + } +}