Something went wrong. Try again.
Caddy module to require at-proto authentication and restrict routes to DIDs
Something went wrong. Try again.
8.2 kB · 259 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260package db
import ( "context" "crypto/rand" "database/sql" "encoding/hex" "encoding/json" "fmt" "os" "path/filepath" "sync/atomic"
"github.com/bluesky-social/indigo/atproto/atcrypto" "github.com/bluesky-social/indigo/atproto/auth/oauth" "github.com/bluesky-social/indigo/atproto/syntax" _ "github.com/mattn/go-sqlite3")
// Ensure DB implements ClientAuthStorevar _ oauth.ClientAuthStore = (*Store)(nil)
// Store handles SQLite persistence for the plugin.type Store struct { db *sql.DB cleanupCounter uint32}
// NewStore initializes a new SQLite-backed storage.func NewStore(path string) (*Store, error) { if dir := filepath.Dir(path); dir != "" && dir != "." { if err := os.MkdirAll(dir, 0755); err != nil { return nil, fmt.Errorf("failed to create directory for database: %w", err) } }
// Enable WAL mode for better concurrency db, err := sql.Open("sqlite3", path+"?_journal_mode=WAL&_busy_timeout=5000") if err != nil { return nil, fmt.Errorf("failed to open database: %w", err) }
if err := db.Ping(); err != nil { return nil, fmt.Errorf("failed to ping database: %w", err) }
store := &Store{db: db} if err := store.initSchema(); err != nil { return nil, fmt.Errorf("failed to initialize schema: %w", err) }
return store, nil}
func (s *Store) initSchema() error { const schema = ` CREATE TABLE IF NOT EXISTS auth_requests ( state TEXT PRIMARY KEY, data TEXT NOT NULL, created_at DATETIME DEFAULT CURRENT_TIMESTAMP );
CREATE TABLE IF NOT EXISTS auth_sessions ( did TEXT NOT NULL, session_id TEXT NOT NULL, data TEXT NOT NULL, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP, PRIMARY KEY (did, session_id) );
CREATE TABLE IF NOT EXISTS system_keys ( id TEXT PRIMARY KEY, key_data BLOB NOT NULL, created_at DATETIME DEFAULT CURRENT_TIMESTAMP ); ` _, err := s.db.Exec(schema) return err}
// Close closes the underlying database.func (s *Store) Close() error { return s.db.Close()}
// GetClientKey retrieves the main client key, generating it if it doesn't exist.// Returns the private key and its ID (a hash of the public key or a random string).func (s *Store) GetClientKey(ctx context.Context) (atcrypto.PrivateKey, string, error) { var keyData []byte err := s.db.QueryRowContext(ctx, "SELECT key_data FROM system_keys WHERE id = 'client_key'").Scan(&keyData) if err == sql.ErrNoRows { // Generate a new P-256 key pk, err := atcrypto.GeneratePrivateKeyP256() if err != nil { return nil, "", fmt.Errorf("failed to generate new client key: %w", err) }
keyData = pk.Bytes()
_, err = s.db.ExecContext(ctx, "INSERT INTO system_keys (id, key_data) VALUES ('client_key', ?)", keyData) if err != nil { return nil, "", fmt.Errorf("failed to save generated client key: %w", err) }
return pk, "client_key", nil } else if err != nil { return nil, "", fmt.Errorf("failed to load client key: %w", err) }
pk, err := atcrypto.ParsePrivateBytesP256(keyData) if err != nil { return nil, "", fmt.Errorf("failed to parse existing client key: %w", err) }
return pk, "client_key", nil}
// GetCookieSecret retrieves or generates the cookie secret.func (s *Store) GetCookieSecret(ctx context.Context) (string, error) { var secret []byte err := s.db.QueryRowContext(ctx, "SELECT key_data FROM system_keys WHERE id = 'cookie_secret'").Scan(&secret) if err == sql.ErrNoRows { // Generate new random 32 byte secret rawSecret := make([]byte, 32) if _, err := rand.Read(rawSecret); err != nil { return "", fmt.Errorf("failed to generate cookie secret: %w", err) } secretStr := hex.EncodeToString(rawSecret)
_, err = s.db.ExecContext(ctx, "INSERT INTO system_keys (id, key_data) VALUES ('cookie_secret', ?)", []byte(secretStr)) if err != nil { return "", fmt.Errorf("failed to save cookie secret: %w", err) } return secretStr, nil } else if err != nil { return "", fmt.Errorf("failed to load cookie secret: %w", err) }
return string(secret), nil}
// GetSession retrieves session data from the database.func (s *Store) GetSession(ctx context.Context, did syntax.DID, sessionID string) (*oauth.ClientSessionData, error) { var dataStr string err := s.db.QueryRowContext(ctx, "SELECT data FROM auth_sessions WHERE did = ? AND session_id = ?", did.String(), sessionID).Scan(&dataStr) if err != nil { if err == sql.ErrNoRows { return nil, nil } return nil, fmt.Errorf("failed to query session: %w", err) }
var sessionData oauth.ClientSessionData if err := json.Unmarshal([]byte(dataStr), &sessionData); err != nil { return nil, fmt.Errorf("failed to parse session data: %w", err) }
return &sessionData, nil}
// SaveSession saves session data into the database.func (s *Store) SaveSession(ctx context.Context, sess oauth.ClientSessionData) error { dataBytes, err := json.Marshal(sess) if err != nil { return fmt.Errorf("failed to serialize session data: %w", err) }
_, err = s.db.ExecContext(ctx, ` INSERT INTO auth_sessions (did, session_id, data, updated_at) VALUES (?, ?, ?, CURRENT_TIMESTAMP) ON CONFLICT(did, session_id) DO UPDATE SET data = excluded.data, updated_at = CURRENT_TIMESTAMP `, sess.AccountDID.String(), sess.SessionID, string(dataBytes)) if err != nil { return fmt.Errorf("failed to save session: %w", err) }
return nil}
// DeleteSession removes session data from the database.func (s *Store) DeleteSession(ctx context.Context, did syntax.DID, sessionID string) error { _, err := s.db.ExecContext(ctx, "DELETE FROM auth_sessions WHERE did = ? AND session_id = ?", did.String(), sessionID) if err != nil { return fmt.Errorf("failed to delete session: %w", err) } return nil}
// Cleanup removes expired auth requests and stale sessions to prevent database bloat.func (s *Store) Cleanup(ctx context.Context) error { // Delete auth_requests older than 1 hour _, err := s.db.ExecContext(ctx, "DELETE FROM auth_requests WHERE created_at < datetime('now', '-1 hour')") if err != nil { return fmt.Errorf("cleanup auth_requests failed: %w", err) }
// Delete auth_sessions older than 30 days _, err = s.db.ExecContext(ctx, "DELETE FROM auth_sessions WHERE updated_at < datetime('now', '-30 days')") if err != nil { return fmt.Errorf("cleanup auth_sessions failed: %w", err) }
return nil}
// GetAuthRequestInfo retrieves the auth request data by state.func (s *Store) GetAuthRequestInfo(ctx context.Context, state string) (*oauth.AuthRequestData, error) { var dataStr string err := s.db.QueryRowContext(ctx, "SELECT data FROM auth_requests WHERE state = ?", state).Scan(&dataStr) if err != nil { if err == sql.ErrNoRows { // AuthRequestData not found return nil, fmt.Errorf("auth request info not found for state") } return nil, fmt.Errorf("failed to query auth request: %w", err) }
var reqData oauth.AuthRequestData if err := json.Unmarshal([]byte(dataStr), &reqData); err != nil { return nil, fmt.Errorf("failed to parse auth request data: %w", err) }
return &reqData, nil}
// SaveAuthRequestInfo saves the auth request info.func (s *Store) SaveAuthRequestInfo(ctx context.Context, info oauth.AuthRequestData) error { dataBytes, err := json.Marshal(info) if err != nil { return fmt.Errorf("failed to serialize auth request data: %w", err) }
// It shouldn't exist, but we can do an INSERT OR REPLACE just in case. _, err = s.db.ExecContext(ctx, ` INSERT OR REPLACE INTO auth_requests (state, data, created_at) VALUES (?, ?, CURRENT_TIMESTAMP) `, info.State, string(dataBytes)) if err != nil { return fmt.Errorf("failed to save auth request: %w", err) }
// Trigger occasional cleanup of stale auth requests and sessions if atomic.AddUint32(&s.cleanupCounter, 1)%100 == 0 { go func() { _ = s.Cleanup(context.Background()) }() }
return nil}
// DeleteAuthRequestInfo removes the auth request info after it's been used or expired.func (s *Store) DeleteAuthRequestInfo(ctx context.Context, state string) error { _, err := s.db.ExecContext(ctx, "DELETE FROM auth_requests WHERE state = ?", state) if err != nil { return fmt.Errorf("failed to delete auth request: %w", err) } return nil}