package pocketbase import ( "context" "database/sql" "errors" "fmt" "strings" "github.com/dariuszjaniak/warp/server/internal/storage" ) // Admin is the PocketBase implementation of storage.AdminStore. type Admin struct{ *base } // registrationTokenColumns reads uses_allowed as NULL for a token without a // limit, as the SQL drivers store it. const registrationTokenColumns = `token, CASE WHEN has_uses_allowed THEN uses_allowed END, completed, expires_at_ms, created_at_ms` // CreateRegistrationToken implements storage.AdminStore. func (s *Admin) CreateRegistrationToken(ctx context.Context, rt storage.RegistrationToken) error { err := s.write(ctx, func(t *tx) error { taken, err := t.exists("admin_registration_tokens", `token = ?`, rt.Token) if err != nil { return err } if taken { return storage.ErrRegistrationTokenExists } row := rec{ "token": rt.Token, "completed": rt.Completed, "expires_at_ms": rt.ExpiresAtMS, "created_at_ms": rt.CreatedAtMS, } setUsesAllowed(row, rt.UsesAllowed) t.put("admin_registration_tokens", row) return nil }) switch { case errors.Is(err, storage.ErrRegistrationTokenExists): return err case err != nil: return fmt.Errorf("create a registration token: %w", err) } return nil } // setUsesAllowed writes an optional use limit into a row. func setUsesAllowed(row rec, n *int) { if n == nil { row["uses_allowed"], row["has_uses_allowed"] = int64(0), false return } row["uses_allowed"], row["has_uses_allowed"] = int64(*n), true } // RegistrationToken implements storage.AdminStore. func (s *Admin) RegistrationToken(ctx context.Context, token string) (storage.RegistrationToken, error) { rt, err := scanRegistrationToken(s.db.QueryRowContext(ctx, `SELECT `+registrationTokenColumns+` FROM admin_registration_tokens WHERE token = ?`, token).Scan) switch { case errors.Is(err, sql.ErrNoRows): return storage.RegistrationToken{}, storage.ErrNoSuchRegistrationToken case err != nil: return storage.RegistrationToken{}, fmt.Errorf("read a registration token: %w", err) } return rt, nil } // RegistrationTokens implements storage.AdminStore. func (s *Admin) RegistrationTokens(ctx context.Context) ([]storage.RegistrationToken, error) { rows, err := s.db.QueryContext(ctx, `SELECT `+registrationTokenColumns+` FROM admin_registration_tokens ORDER BY token`) if err != nil { return nil, fmt.Errorf("list registration tokens: %w", err) } defer func() { _ = rows.Close() }() out := []storage.RegistrationToken{} for rows.Next() { rt, err := scanRegistrationToken(rows.Scan) if err != nil { return nil, fmt.Errorf("list registration tokens: %w", err) } out = append(out, rt) } return out, rows.Err() } // UpdateRegistrationToken implements storage.AdminStore. func (s *Admin) UpdateRegistrationToken(ctx context.Context, token string, usesAllowed *int, expiresAtMS int64) error { err := s.write(ctx, func(t *tx) error { row := rec{"expires_at_ms": expiresAtMS} setUsesAllowed(row, usesAllowed) n, err := t.setWhere("admin_registration_tokens", row, `token = ?`, token) if err != nil { return fmt.Errorf("update a registration token: %w", err) } if n == 0 { return storage.ErrNoSuchRegistrationToken } return nil }) return err } // DeleteRegistrationToken implements storage.AdminStore. func (s *Admin) DeleteRegistrationToken(ctx context.Context, token string) error { return s.write(ctx, func(t *tx) error { found, err := t.delKey("admin_registration_tokens", token) if err != nil { return fmt.Errorf("delete a registration token: %w", err) } if !found { return storage.ErrNoSuchRegistrationToken } return nil }) } // UseRegistrationToken implements storage.AdminStore. The check and the // spend happen under the write lock, so two registrations cannot both take // the last use. func (s *Admin) UseRegistrationToken(ctx context.Context, token string, nowMS int64) error { return s.write(ctx, func(t *tx) error { var completed int64 err := t.queryRow(`SELECT completed FROM admin_registration_tokens WHERE token = ? AND (NOT has_uses_allowed OR completed < uses_allowed) AND (expires_at_ms = 0 OR expires_at_ms > ?)`, token, nowMS).Scan(&completed) switch { case errors.Is(err, sql.ErrNoRows): return storage.ErrNoSuchRegistrationToken case err != nil: return fmt.Errorf("spend a registration token: %w", err) } t.put("admin_registration_tokens", rec{"token": token, "completed": completed + 1}) return nil }) } // ReleaseRegistrationToken implements storage.AdminStore. A token deleted in // the meantime has nothing to give back to, which is not an error. func (s *Admin) ReleaseRegistrationToken(ctx context.Context, token string) error { err := s.write(ctx, func(t *tx) error { var completed int64 err := t.queryRow(`SELECT completed FROM admin_registration_tokens WHERE token = ? AND completed > 0`, token). Scan(&completed) switch { case errors.Is(err, sql.ErrNoRows): return nil case err != nil: return err } t.put("admin_registration_tokens", rec{"token": token, "completed": completed - 1}) return nil }) if err != nil { return fmt.Errorf("release a registration token: %w", err) } return nil } func scanRegistrationToken(scan func(...any) error) (storage.RegistrationToken, error) { var ( rt storage.RegistrationToken uses sql.NullInt64 ) if err := scan(&rt.Token, &uses, &rt.Completed, &rt.ExpiresAtMS, &rt.CreatedAtMS); err != nil { return storage.RegistrationToken{}, err } if uses.Valid { n := int(uses.Int64) rt.UsesAllowed = &n } return rt, nil } const ( auditColumns = `audit_id, ts_ms, actor, action, target, details_json` defaultAuditPage = 50 maxAuditPage = 1000 ) // AppendAudit implements storage.AdminStore. func (s *Admin) AppendAudit(ctx context.Context, e storage.AuditEntry) (int64, error) { details := string(e.Details) if details == "" { details = "{}" } var id int64 err := s.write(ctx, func(t *tx) error { var err error if id, err = t.next(adminSequences, seqAudit); err != nil { return err } t.put("admin_audit_log", rec{ "audit_id": id, "ts_ms": e.TSMS, "actor": e.Actor, "action": e.Action, "target": e.Target, "details_json": details, }) return nil }) if err != nil { return 0, fmt.Errorf("append to the audit log: %w", err) } return id, nil } // AuditLog implements storage.AdminStore. func (s *Admin) AuditLog(ctx context.Context, f storage.AuditFilter) ([]storage.AuditEntry, error) { var ( conditions []string args []any ) if f.Before > 0 { conditions = append(conditions, `audit_id < ?`) args = append(args, f.Before) } if f.Actor != "" { conditions = append(conditions, `actor = ?`) args = append(args, f.Actor) } if f.Action != "" { conditions = append(conditions, `action = ?`) args = append(args, f.Action) } limit := f.Limit switch { case limit <= 0: limit = defaultAuditPage case limit > maxAuditPage: limit = maxAuditPage } args = append(args, limit) query := `SELECT ` + auditColumns + ` FROM admin_audit_log` if len(conditions) > 0 { query += ` WHERE ` + strings.Join(conditions, " AND ") } query += ` ORDER BY audit_id DESC LIMIT ?` rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { return nil, fmt.Errorf("read the audit log: %w", err) } defer func() { _ = rows.Close() }() out := []storage.AuditEntry{} for rows.Next() { var ( e storage.AuditEntry details string ) if err := rows.Scan(&e.ID, &e.TSMS, &e.Actor, &e.Action, &e.Target, &details); err != nil { return nil, fmt.Errorf("read the audit log: %w", err) } e.Details = []byte(details) out = append(out, e) } return out, rows.Err() } // ServerNoticeRoom implements storage.AdminStore. func (s *Admin) ServerNoticeRoom(ctx context.Context, localpart string) (string, error) { var roomID string err := s.db.QueryRowContext(ctx, `SELECT room_id FROM admin_server_notice_rooms WHERE localpart = ?`, localpart). Scan(&roomID) switch { case errors.Is(err, sql.ErrNoRows): return "", storage.ErrNoServerNoticeRoom case err != nil: return "", fmt.Errorf("read the server notice room of %q: %w", localpart, err) } return roomID, nil } // SetServerNoticeRoom implements storage.AdminStore. func (s *Admin) SetServerNoticeRoom(ctx context.Context, localpart, roomID string) error { err := s.write(ctx, func(t *tx) error { t.put("admin_server_notice_rooms", rec{"localpart": localpart, "room_id": roomID}) return nil }) if err != nil { return fmt.Errorf("record the server notice room of %q: %w", localpart, err) } return nil } // AppServices is the PocketBase implementation of storage.AppServiceStore. type AppServices struct{ *base } // AppServiceState implements storage.AppServiceStore. func (s *AppServices) AppServiceState(ctx context.Context, id string, startPos int64) (storage.AppServiceState, error) { var st storage.AppServiceState err := s.write(ctx, func(t *tx) error { err := t.queryRow(`SELECT as_id, stream_pos, txn_id, txn_to FROM appservice_state WHERE as_id = ?`, id). Scan(&st.ID, &st.StreamPos, &st.TxnID, &st.TxnTo) switch { case errors.Is(err, sql.ErrNoRows): st = storage.AppServiceState{ID: id, StreamPos: startPos} t.put("appservice_state", rec{"as_id": id, "stream_pos": startPos, "txn_id": int64(0), "txn_to": int64(0)}) return nil case err != nil: return err } return nil }) if err != nil { return storage.AppServiceState{}, fmt.Errorf("read the cursor of application service %q: %w", id, err) } return st, nil } // BeginTransaction implements storage.AppServiceStore. func (s *AppServices) BeginTransaction(ctx context.Context, id string, txnID, to int64) error { err := s.write(ctx, func(t *tx) error { _, err := t.setWhere("appservice_state", rec{"txn_id": txnID, "txn_to": to}, `as_id = ?`, id) return err }) if err != nil { return fmt.Errorf("begin transaction %d of application service %q: %w", txnID, id, err) } return nil } // CompleteTransaction implements storage.AppServiceStore. func (s *AppServices) CompleteTransaction(ctx context.Context, id string, to int64) error { err := s.write(ctx, func(t *tx) error { _, err := t.setWhere("appservice_state", rec{"stream_pos": to, "txn_to": int64(0)}, `as_id = ?`, id) return err }) if err != nil { return fmt.Errorf("advance the cursor of application service %q: %w", id, err) } return nil } var ( _ storage.AdminStore = (*Admin)(nil) _ storage.AppServiceStore = (*AppServices)(nil) )