diff --git a/internal/oauth/crypto.go b/internal/oauth/crypto.go new file mode 100644 index 0000000..e916478 --- /dev/null +++ b/internal/oauth/crypto.go @@ -0,0 +1,57 @@ +package oauth + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "errors" + "fmt" + "io" +) + +const KeySize = 32 + +var ErrInvalidKey = errors.New("oauth: key must be 32 bytes") + +func Encrypt(plaintext, key []byte) ([]byte, error) { + if len(key) != KeySize { + return nil, ErrInvalidKey + } + block, err := aes.NewCipher(key) + if err != nil { + return nil, fmt.Errorf("aes: %w", err) + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, fmt.Errorf("gcm: %w", err) + } + nonce := make([]byte, gcm.NonceSize()) + if _, err := io.ReadFull(rand.Reader, nonce); err != nil { + return nil, fmt.Errorf("nonce: %w", err) + } + return gcm.Seal(nonce, nonce, plaintext, nil), nil +} + +func Decrypt(ciphertext, key []byte) ([]byte, error) { + if len(key) != KeySize { + return nil, ErrInvalidKey + } + block, err := aes.NewCipher(key) + if err != nil { + return nil, fmt.Errorf("aes: %w", err) + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, fmt.Errorf("gcm: %w", err) + } + ns := gcm.NonceSize() + if len(ciphertext) < ns { + return nil, errors.New("oauth: ciphertext too short") + } + nonce, ct := ciphertext[:ns], ciphertext[ns:] + pt, err := gcm.Open(nil, nonce, ct, nil) + if err != nil { + return nil, fmt.Errorf("open: %w", err) + } + return pt, nil +} diff --git a/internal/oauth/crypto_test.go b/internal/oauth/crypto_test.go new file mode 100644 index 0000000..65c0468 --- /dev/null +++ b/internal/oauth/crypto_test.go @@ -0,0 +1,63 @@ +package oauth + +import ( + "bytes" + "crypto/rand" + "testing" +) + +func newKey(t *testing.T) []byte { + t.Helper() + k := make([]byte, 32) + if _, err := rand.Read(k); err != nil { + t.Fatal(err) + } + return k +} + +func TestEncryptDecrypt_RoundTrip(t *testing.T) { + key := newKey(t) + plaintext := []byte("hello-oauth-token") + + ct, err := Encrypt(plaintext, key) + if err != nil { + t.Fatalf("encrypt: %v", err) + } + if bytes.Equal(ct, plaintext) { + t.Fatal("ciphertext == plaintext") + } + + pt, err := Decrypt(ct, key) + if err != nil { + t.Fatalf("decrypt: %v", err) + } + if !bytes.Equal(pt, plaintext) { + t.Fatalf("got %q, want %q", pt, plaintext) + } +} + +func TestEncrypt_RejectsKeyOfWrongLength(t *testing.T) { + _, err := Encrypt([]byte("x"), make([]byte, 16)) + if err == nil { + t.Fatal("expected error for 16-byte key") + } +} + +func TestDecrypt_FailsOnTamperedCiphertext(t *testing.T) { + key := newKey(t) + ct, _ := Encrypt([]byte("tamper-me"), key) + ct[len(ct)-1] ^= 0x01 + + if _, err := Decrypt(ct, key); err == nil { + t.Fatal("expected error for tampered ciphertext") + } +} + +func TestEncrypt_NonDeterministic(t *testing.T) { + key := newKey(t) + a, _ := Encrypt([]byte("same"), key) + b, _ := Encrypt([]byte("same"), key) + if bytes.Equal(a, b) { + t.Fatal("two encryptions of same plaintext produced same ciphertext (nonce reuse)") + } +} diff --git a/internal/store/migrate_integration_test.go b/internal/store/migrate_integration_test.go index 56bd5ae..5ba77a2 100644 --- a/internal/store/migrate_integration_test.go +++ b/internal/store/migrate_integration_test.go @@ -10,6 +10,42 @@ import ( "github.com/smallchungus/disttaskqueue/internal/testutil" ) +func TestMigrate_AppliesAllPhase3ATables(t *testing.T) { + pool := testutil.StartPostgres(t) + if err := store.Migrate(context.Background(), pool.Config().ConnString()); err != nil { + t.Fatalf("migrate: %v", err) + } + + wantTables := []string{"users", "oauth_tokens", "gmail_sync_state", "processed_emails", "pipeline_jobs", "job_status_history"} + for _, table := range wantTables { + var n int + err := pool.QueryRow(context.Background(), + `SELECT count(*) FROM information_schema.tables WHERE table_schema='public' AND table_name=$1`, table, + ).Scan(&n) + if err != nil { + t.Fatal(err) + } + if n != 1 { + t.Fatalf("table %s missing", table) + } + } + + wantCols := []string{"user_id", "gmail_message_id", "is_synthetic"} + for _, col := range wantCols { + var n int + err := pool.QueryRow(context.Background(), + `SELECT count(*) FROM information_schema.columns + WHERE table_schema='public' AND table_name='pipeline_jobs' AND column_name=$1`, col, + ).Scan(&n) + if err != nil { + t.Fatal(err) + } + if n != 1 { + t.Fatalf("pipeline_jobs.%s missing", col) + } + } +} + func TestMigrate_AppliesInitialSchema(t *testing.T) { pool := testutil.StartPostgres(t) diff --git a/internal/store/migrations/0002_users_oauth_sync_processed.down.sql b/internal/store/migrations/0002_users_oauth_sync_processed.down.sql new file mode 100644 index 0000000..ffa77ca --- /dev/null +++ b/internal/store/migrations/0002_users_oauth_sync_processed.down.sql @@ -0,0 +1,12 @@ +DROP INDEX IF EXISTS pipeline_jobs_user_message_unique; +DROP INDEX IF EXISTS pipeline_jobs_user_stage_status_idx; + +ALTER TABLE pipeline_jobs + DROP COLUMN IF EXISTS is_synthetic, + DROP COLUMN IF EXISTS gmail_message_id, + DROP COLUMN IF EXISTS user_id; + +DROP TABLE IF EXISTS processed_emails; +DROP TABLE IF EXISTS gmail_sync_state; +DROP TABLE IF EXISTS oauth_tokens; +DROP TABLE IF EXISTS users; diff --git a/internal/store/migrations/0002_users_oauth_sync_processed.up.sql b/internal/store/migrations/0002_users_oauth_sync_processed.up.sql new file mode 100644 index 0000000..1e31f3f --- /dev/null +++ b/internal/store/migrations/0002_users_oauth_sync_processed.up.sql @@ -0,0 +1,40 @@ +CREATE TABLE users ( + id UUID PRIMARY KEY, + email TEXT NOT NULL UNIQUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE TABLE oauth_tokens ( + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + provider TEXT NOT NULL, + access_ct BYTEA NOT NULL, + refresh_ct BYTEA NOT NULL, + expires_at TIMESTAMPTZ NOT NULL, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (user_id, provider) +); + +CREATE TABLE gmail_sync_state ( + user_id UUID PRIMARY KEY REFERENCES users(id) ON DELETE CASCADE, + history_id TEXT NOT NULL, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE TABLE processed_emails ( + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + gmail_message_id TEXT NOT NULL, + drive_folder_id TEXT NOT NULL, + processed_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (user_id, gmail_message_id) +); + +ALTER TABLE pipeline_jobs + ADD COLUMN user_id UUID NULL REFERENCES users(id) ON DELETE SET NULL, + ADD COLUMN gmail_message_id TEXT NULL, + ADD COLUMN is_synthetic BOOLEAN NOT NULL DEFAULT false; + +CREATE INDEX pipeline_jobs_user_stage_status_idx ON pipeline_jobs (user_id, stage, status); + +CREATE UNIQUE INDEX pipeline_jobs_user_message_unique + ON pipeline_jobs (user_id, gmail_message_id) + WHERE status NOT IN ('done', 'dead') AND is_synthetic = false AND user_id IS NOT NULL AND gmail_message_id IS NOT NULL; diff --git a/internal/store/oauth.go b/internal/store/oauth.go new file mode 100644 index 0000000..7d10645 --- /dev/null +++ b/internal/store/oauth.go @@ -0,0 +1,46 @@ +package store + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" +) + +var ErrOAuthTokenNotFound = errors.New("oauth token not found") + +func (s *Store) SaveOAuthToken(ctx context.Context, t OAuthToken) error { + const q = ` + INSERT INTO oauth_tokens (user_id, provider, access_ct, refresh_ct, expires_at, updated_at) + VALUES ($1, $2, $3, $4, $5, now()) + ON CONFLICT (user_id, provider) DO UPDATE + SET access_ct = EXCLUDED.access_ct, + refresh_ct = EXCLUDED.refresh_ct, + expires_at = EXCLUDED.expires_at, + updated_at = now()` + + if _, err := s.pool.Exec(ctx, q, t.UserID, t.Provider, t.AccessCT, t.RefreshCT, t.ExpiresAt); err != nil { + return fmt.Errorf("save oauth: %w", err) + } + return nil +} + +func (s *Store) GetOAuthToken(ctx context.Context, userID uuid.UUID, provider string) (OAuthToken, error) { + const q = ` + SELECT user_id, provider, access_ct, refresh_ct, expires_at, updated_at + FROM oauth_tokens + WHERE user_id = $1 AND provider = $2` + + var t OAuthToken + err := s.pool.QueryRow(ctx, q, userID, provider).Scan( + &t.UserID, &t.Provider, &t.AccessCT, &t.RefreshCT, &t.ExpiresAt, &t.UpdatedAt) + if errors.Is(err, pgx.ErrNoRows) { + return OAuthToken{}, ErrOAuthTokenNotFound + } + if err != nil { + return OAuthToken{}, fmt.Errorf("get oauth: %w", err) + } + return t, nil +} diff --git a/internal/store/store_integration_test.go b/internal/store/store_integration_test.go index 43b08a8..b52f985 100644 --- a/internal/store/store_integration_test.go +++ b/internal/store/store_integration_test.go @@ -198,6 +198,101 @@ func TestListReadyRetryJobs_ReturnsQueuedWithLastErrorAndDueNextRun(t *testing.T _, _ = j2, j3 } +func TestCreateUser_PersistsRow(t *testing.T) { + s := newStore(t) + ctx := context.Background() + + u, err := s.CreateUser(ctx, "test@example.com") + if err != nil { + t.Fatalf("create: %v", err) + } + if u.ID.String() == "" || u.Email != "test@example.com" { + t.Fatalf("user: %+v", u) + } + + got, err := s.GetUserByEmail(ctx, "test@example.com") + if err != nil { + t.Fatalf("get: %v", err) + } + if got.ID != u.ID { + t.Fatalf("id mismatch: %s vs %s", got.ID, u.ID) + } +} + +func TestGetUserByEmail_ReturnsErrUserNotFound(t *testing.T) { + s := newStore(t) + _, err := s.GetUserByEmail(context.Background(), "nobody@example.com") + if !errors.Is(err, store.ErrUserNotFound) { + t.Fatalf("got %v, want ErrUserNotFound", err) + } +} + +func TestSaveOAuthToken_RoundTrip(t *testing.T) { + s := newStore(t) + ctx := context.Background() + u, _ := s.CreateUser(ctx, "oauth@example.com") + + expires := time.Now().Add(1 * time.Hour).UTC().Truncate(time.Second) + in := store.OAuthToken{ + UserID: u.ID, + Provider: "google", + AccessCT: []byte("encrypted-access"), + RefreshCT: []byte("encrypted-refresh"), + ExpiresAt: expires, + } + if err := s.SaveOAuthToken(ctx, in); err != nil { + t.Fatalf("save: %v", err) + } + + got, err := s.GetOAuthToken(ctx, u.ID, "google") + if err != nil { + t.Fatalf("get: %v", err) + } + if got.UserID != in.UserID || got.Provider != in.Provider { + t.Fatalf("identity mismatch: %+v", got) + } + if string(got.AccessCT) != string(in.AccessCT) || string(got.RefreshCT) != string(in.RefreshCT) { + t.Fatalf("ciphertext mismatch") + } + if !got.ExpiresAt.Equal(expires) { + t.Fatalf("expires: got %v want %v", got.ExpiresAt, expires) + } +} + +func TestSaveOAuthToken_UpsertsExistingRow(t *testing.T) { + s := newStore(t) + ctx := context.Background() + u, _ := s.CreateUser(ctx, "upsert@example.com") + + first := store.OAuthToken{UserID: u.ID, Provider: "google", + AccessCT: []byte("v1-access"), RefreshCT: []byte("v1-refresh"), + ExpiresAt: time.Now().Add(time.Hour)} + _ = s.SaveOAuthToken(ctx, first) + + second := store.OAuthToken{UserID: u.ID, Provider: "google", + AccessCT: []byte("v2-access"), RefreshCT: []byte("v2-refresh"), + ExpiresAt: time.Now().Add(2 * time.Hour)} + if err := s.SaveOAuthToken(ctx, second); err != nil { + t.Fatalf("upsert: %v", err) + } + + got, _ := s.GetOAuthToken(ctx, u.ID, "google") + if string(got.AccessCT) != "v2-access" { + t.Fatalf("not upserted: %s", got.AccessCT) + } +} + +func TestGetOAuthToken_ReturnsErrOAuthTokenNotFound(t *testing.T) { + s := newStore(t) + ctx := context.Background() + u, _ := s.CreateUser(ctx, "missing@example.com") + + _, err := s.GetOAuthToken(ctx, u.ID, "google") + if !errors.Is(err, store.ErrOAuthTokenNotFound) { + t.Fatalf("got %v, want ErrOAuthTokenNotFound", err) + } +} + func jobIDs(js []store.Job) []string { out := make([]string, len(js)) for i, j := range js { diff --git a/internal/store/types.go b/internal/store/types.go index 0eb3ccd..50d50ea 100644 --- a/internal/store/types.go +++ b/internal/store/types.go @@ -36,3 +36,18 @@ type NewJob struct { Stage string Payload json.RawMessage } + +type User struct { + ID uuid.UUID + Email string + CreatedAt time.Time +} + +type OAuthToken struct { + UserID uuid.UUID + Provider string + AccessCT []byte + RefreshCT []byte + ExpiresAt time.Time + UpdatedAt time.Time +} diff --git a/internal/store/users.go b/internal/store/users.go new file mode 100644 index 0000000..e8488a7 --- /dev/null +++ b/internal/store/users.go @@ -0,0 +1,39 @@ +package store + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" +) + +var ErrUserNotFound = errors.New("user not found") + +func (s *Store) CreateUser(ctx context.Context, email string) (User, error) { + id := uuid.New() + const q = ` + INSERT INTO users (id, email) VALUES ($1, $2) + RETURNING id, email, created_at` + + var u User + if err := s.pool.QueryRow(ctx, q, id, email).Scan(&u.ID, &u.Email, &u.CreatedAt); err != nil { + return User{}, fmt.Errorf("create user: %w", err) + } + return u, nil +} + +func (s *Store) GetUserByEmail(ctx context.Context, email string) (User, error) { + const q = `SELECT id, email, created_at FROM users WHERE email = $1` + + var u User + err := s.pool.QueryRow(ctx, q, email).Scan(&u.ID, &u.Email, &u.CreatedAt) + if errors.Is(err, pgx.ErrNoRows) { + return User{}, ErrUserNotFound + } + if err != nil { + return User{}, fmt.Errorf("get user: %w", err) + } + return u, nil +}