(READ ONLY) Margin is an open annotation layer for the internet. Powered by the AT Protocol.
Something went wrong. Try again.
Go
at ui-refactor
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205package api
import ( "bytes" "context" "crypto/ecdsa" "crypto/x509" "encoding/pem" "fmt" "log" "net/http" "os" "time"
"margin.at/internal/db" "margin.at/internal/oauth" "margin.at/internal/xrpc")
type TokenRefresher struct { db *db.DB privateKey *ecdsa.PrivateKey baseURL string}
func NewTokenRefresher(database *db.DB, privateKey *ecdsa.PrivateKey) *TokenRefresher { return &TokenRefresher{ db: database, privateKey: privateKey, baseURL: os.Getenv("BASE_URL"), }}
func (tr *TokenRefresher) getOAuthClient(r *http.Request) *oauth.Client { baseURL := tr.baseURL if baseURL == "" { scheme := "http" if r.TLS != nil || r.Header.Get("X-Forwarded-Proto") == "https" { scheme = "https" } baseURL = fmt.Sprintf("%s://%s", scheme, r.Host) }
if len(baseURL) > 0 && baseURL[len(baseURL)-1] == '/' { baseURL = baseURL[:len(baseURL)-1] }
clientID := baseURL + "/client-metadata.json" redirectURI := baseURL + "/auth/callback"
return oauth.NewClient(clientID, redirectURI, tr.privateKey)}
type SessionData struct { ID string DID string Handle string AccessToken string RefreshToken string DPoPKey *ecdsa.PrivateKey PDS string}
func (tr *TokenRefresher) GetSessionWithAutoRefresh(r *http.Request) (*SessionData, error) { sessionID := ""
cookie, err := r.Cookie("margin_session") if err == nil { sessionID = cookie.Value } else { sessionID = r.Header.Get("X-Session-Token") }
if sessionID == "" { return nil, fmt.Errorf("not authenticated") }
did, handle, accessToken, refreshToken, dpopKeyStr, err := tr.db.GetSession(sessionID) if err != nil { return nil, fmt.Errorf("session expired") }
block, _ := pem.Decode([]byte(dpopKeyStr)) if block == nil { return nil, fmt.Errorf("invalid session DPoP key") } dpopKey, err := x509.ParseECPrivateKey(block.Bytes) if err != nil { return nil, fmt.Errorf("invalid session DPoP key") }
pds, err := xrpc.ResolveDIDToPDS(did) if err != nil { return nil, fmt.Errorf("failed to resolve PDS") }
return &SessionData{ ID: sessionID, DID: did, Handle: handle, AccessToken: accessToken, RefreshToken: refreshToken, DPoPKey: dpopKey, PDS: pds, }, nil}
func (tr *TokenRefresher) RefreshSessionToken(r *http.Request, session *SessionData) (*SessionData, error) { if session.ID == "" { return nil, fmt.Errorf("invalid session ID") }
oauthClient := tr.getOAuthClient(r) ctx := context.Background()
meta, err := oauthClient.GetAuthServerMetadata(ctx, session.PDS) if err != nil { return nil, fmt.Errorf("failed to get auth server metadata: %w", err) }
tokenResp, _, err := oauthClient.RefreshToken(meta, session.RefreshToken, session.DPoPKey, "") if err != nil { return nil, fmt.Errorf("failed to refresh token: %w", err) }
dpopKeyBytes, err := x509.MarshalECPrivateKey(session.DPoPKey) if err != nil { return nil, fmt.Errorf("failed to marshal DPoP key: %w", err) } dpopKeyPEM := pem.EncodeToMemory(&pem.Block{ Type: "EC PRIVATE KEY", Bytes: dpopKeyBytes, })
newRefreshToken := tokenResp.RefreshToken if newRefreshToken == "" { newRefreshToken = session.RefreshToken }
expiresAt := time.Now().Add(7 * 24 * time.Hour) if err := tr.db.SaveSession( session.ID, session.DID, session.Handle, tokenResp.AccessToken, newRefreshToken, string(dpopKeyPEM), expiresAt, ); err != nil { return nil, fmt.Errorf("failed to save refreshed session: %w", err) }
log.Printf("Successfully refreshed token for user %s", session.Handle)
return &SessionData{ ID: session.ID, DID: session.DID, Handle: session.Handle, AccessToken: tokenResp.AccessToken, RefreshToken: newRefreshToken, DPoPKey: session.DPoPKey, PDS: session.PDS, }, nil}
func IsTokenExpiredError(err error) bool { if err == nil { return false } errStr := err.Error() return bytes.Contains([]byte(errStr), []byte("invalid_token")) && bytes.Contains([]byte(errStr), []byte("exp"))}
func (tr *TokenRefresher) ExecuteWithAutoRefresh( r *http.Request, session *SessionData, fn func(client *xrpc.Client, did string) error,) error { client := xrpc.NewClient(session.PDS, session.AccessToken, session.DPoPKey)
err := fn(client, session.DID) if err == nil { return nil }
if !IsTokenExpiredError(err) { return err }
log.Printf("Token expired for user %s, attempting refresh...", session.Handle)
newSession, refreshErr := tr.RefreshSessionToken(r, session) if refreshErr != nil { return fmt.Errorf("original error: %w; refresh failed: %v", err, refreshErr) }
client = xrpc.NewClient(newSession.PDS, newSession.AccessToken, newSession.DPoPKey) return fn(client, newSession.DID)}
func (tr *TokenRefresher) CreateClientFromSession(session *SessionData) *xrpc.Client { return xrpc.NewClient(session.PDS, session.AccessToken, session.DPoPKey)}