package api
import (
"context"
cryptorand "crypto/rand"
"encoding/base64"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"path/filepath"
"strings"
"sync"
"testing"
"time"
"github.com/bluesky-social/indigo/atproto/atcrypto"
"github.com/bluesky-social/indigo/atproto/auth"
indigooauth "github.com/bluesky-social/indigo/atproto/auth/oauth"
"github.com/bluesky-social/indigo/atproto/identity"
"github.com/bluesky-social/indigo/atproto/syntax"
"tangled.org/core/api/org_tangled"
"tangled.org/core/log"
"tangled.org/core/migrator/config"
"tangled.org/core/migrator/db"
"tangled.org/core/migrator/oauth"
"tangled.org/core/xrpc/serviceauth"
)
type grantStub struct {
mu sync.Mutex
server *httptest.Server
dead bool
refreshes int
requests int
usedRefresh map[string]bool
latestAccess string
mintCalls []string
userDID syntax.DID
userKey atcrypto.PrivateKey
}
func newGrantStub(t *testing.T, userDID syntax.DID) *grantStub {
t.Helper()
userKey, err := atcrypto.GeneratePrivateKeyP256()
if err != nil {
t.Fatal(err)
}
g := &grantStub{usedRefresh: map[string]bool{}, userDID: userDID, userKey: userKey}
g.server = httptest.NewServer(http.HandlerFunc(g.serve))
t.Cleanup(g.server.Close)
return g
}
func (g *grantStub) serve(w http.ResponseWriter, r *http.Request) {
g.mu.Lock()
g.requests++
g.mu.Unlock()
switch r.URL.Path {
case "/oauth/token":
g.token(w, r)
case "/xrpc/com.atproto.server.getServiceAuth":
g.serviceAuth(w, r)
default:
http.NotFound(w, r)
}
}
func (g *grantStub) token(w http.ResponseWriter, r *http.Request) {
if err := r.ParseForm(); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
g.mu.Lock()
defer g.mu.Unlock()
if g.dead || g.usedRefresh[r.Form.Get("refresh_token")] {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadRequest)
_, _ = w.Write([]byte(`{"error":"invalid_grant","error_description":"Refresh token replayed"}`))
return
}
g.usedRefresh[r.Form.Get("refresh_token")] = true
g.refreshes++
g.latestAccess = fmt.Sprintf("pds-oauth-refreshed-%d", g.refreshes)
_ = json.NewEncoder(w).Encode(indigooauth.TokenResponse{
Subject: g.userDID.String(),
Scope: "atproto",
AccessToken: g.latestAccess,
RefreshToken: fmt.Sprintf("refresh-next-%d", g.refreshes),
})
}
func (g *grantStub) serviceAuth(w http.ResponseWriter, r *http.Request) {
g.mu.Lock()
if bearer := strings.TrimPrefix(r.Header.Get("Authorization"), "DPoP "); bearer != g.latestAccess {
g.mu.Unlock()
w.Header().Set("WWW-Authenticate", `Bearer error="invalid_token", error_description="The access token expired"`)
w.WriteHeader(http.StatusUnauthorized)
return
}
g.mintCalls = append(g.mintCalls, r.URL.Query().Get("lxm")+" aud="+r.URL.Query().Get("aud"))
g.mu.Unlock()
lxm, err := syntax.ParseNSID(r.URL.Query().Get("lxm"))
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
token, err := auth.SignServiceAuth(g.userDID, r.URL.Query().Get("aud"), time.Minute, &lxm, g.userKey)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
_ = json.NewEncoder(w).Encode(map[string]string{"token": token})
}
func serverWithGrants(t *testing.T, g *grantStub) (*testServer, string) {
t.Helper()
dir := t.TempDir()
database, err := db.Make(context.Background(), filepath.Join(dir, "api_grant.db"))
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { database.Close() })
rawKey := make([]byte, 32)
_, _ = cryptorand.Read(rawKey)
serviceKey, err := atcrypto.GeneratePrivateKeyP256()
if err != nil {
t.Fatal(err)
}
cfg := &config.Config{
Hostname: "migrator.example.com", PrivateKey: serviceKey.Multibase(),
MasterKey: base64.StdEncoding.EncodeToString(rawKey), WorkDir: dir,
// the stub pds listens on loopback
AllowPrivateNetworks: true,
}
if err := cfg.Validate(); err != nil {
t.Fatal(err)
}
oauthClient, err := oauth.NewClient(cfg, database, identity.NewMockDirectory(), log.New("api-grant-test"))
if err != nil {
t.Fatal(err)
}
store := oauth.NewStore(database, cfg.ParsedMasterKey)
dpopKey, err := atcrypto.GeneratePrivateKeyP256()
if err != nil {
t.Fatal(err)
}
owner := syntax.DID(alice)
if err := store.SaveSession(context.Background(), indigooauth.ClientSessionData{
AccountDID: owner, SessionID: "session", HostURL: g.server.URL, AuthServerURL: g.server.URL,
AuthServerTokenEndpoint: g.server.URL + "/oauth/token",
Scopes: []string{
"atproto",
"rpc:org.tangled.repo.createRepo?aud=*",
"repo:org.tangled.repo.declaration?action=create",
},
AccessToken: "pds-oauth-old", RefreshToken: "refresh-old", DPoPPrivateKeyMultibase: dpopKey.Multibase(),
}); err != nil {
t.Fatal(err)
}
logger := log.New("test")
directory := identity.NewMockDirectory()
priv, err := atcrypto.GeneratePrivateKeyP256()
if err != nil {
t.Fatal(err)
}
pub, err := priv.PublicKey()
if err != nil {
t.Fatal(err)
}
directory.Insert(identity.Identity{
DID: owner,
Keys: map[string]identity.VerificationMethod{
"atproto": {Type: "Multikey", PublicKeyMultibase: pub.Multibase()},
},
})
serviceAuth := serviceauth.NewServiceAuth(logger, directory, cfg.ServiceDid.String())
return &testServer{
server: NewServer(database, cfg, nil, serviceAuth, oauthClient, logger, nil, ""),
db: database,
cfg: cfg,
alice: caller{did: owner, signer: priv},
}, cfg.ServiceDid.String()
}
func TestCreateTaskWithADeadGrantRequiresReauthorization(t *testing.T) {
g := newGrantStub(t, syntax.DID(alice))
g.dead = true
ts, _ := serverWithGrants(t, g)
path := "/xrpc/" + org_tangled.TempMigratorCreateTaskNSID
body := &org_tangled.TempMigratorCreateTask_Input{RequestId: "req-dead", Jobs: []*org_tangled.TempMigratorDefs_NewJob{job("one", false)}}
rec := ts.do(t, http.MethodPost, path, ts.token(t, ts.alice, org_tangled.TempMigratorCreateTaskNSID), body)
if rec.Code != 428 {
t.Fatalf("want 428 Precondition Required, got %d (body: %s)", rec.Code, rec.Body.String())
}
if !strings.Contains(rec.Body.String(), "GrantRequired") {
t.Fatalf("body = %s, want the GrantRequired tag", rec.Body.String())
}
t.Logf("428 body: %s", rec.Body.String())
var batches, jobs int
if err := ts.db.QueryRow("select count(*) from batches").Scan(&batches); err != nil {
t.Fatal(err)
}
if err := ts.db.QueryRow("select count(*) from jobs").Scan(&jobs); err != nil {
t.Fatal(err)
}
t.Logf("rows after dead-grant createTask: batches=%d jobs=%d", batches, jobs)
if batches != 0 || jobs != 0 {
t.Fatalf("a dead grant created %d batches and %d jobs", batches, jobs)
}
g.mu.Lock()
before := g.requests
g.mu.Unlock()
repeat := ts.do(t, http.MethodPost, path, ts.token(t, ts.alice, org_tangled.TempMigratorCreateTaskNSID), body)
if repeat.Code != 428 || !strings.Contains(repeat.Body.String(), "GrantRequired") {
t.Fatalf("repeat createTask = %d %s, want the same 428 GrantRequired", repeat.Code, repeat.Body.String())
}
g.mu.Lock()
defer g.mu.Unlock()
if g.requests != before {
t.Fatalf("repeat createTask made %d stub requests, want 0 (the dead row must be gone)", g.requests-before)
}
}
func TestCreateTaskWithAHealthyGrantStillCreatesTheTask(t *testing.T) {
g := newGrantStub(t, syntax.DID(alice))
ts, serviceDid := serverWithGrants(t, g)
path := "/xrpc/" + org_tangled.TempMigratorCreateTaskNSID
body := &org_tangled.TempMigratorCreateTask_Input{RequestId: "req-alive", Jobs: []*org_tangled.TempMigratorDefs_NewJob{job("one", false), job("two", false)}}
rec := ts.do(t, http.MethodPost, path, ts.token(t, ts.alice, org_tangled.TempMigratorCreateTaskNSID), body)
if rec.Code != http.StatusAccepted {
t.Fatalf("want 202, got %d (body: %s)", rec.Code, rec.Body.String())
}
var batches, jobs int
if err := ts.db.QueryRow("select count(*) from batches").Scan(&batches); err != nil {
t.Fatal(err)
}
if err := ts.db.QueryRow("select count(*) from jobs").Scan(&jobs); err != nil {
t.Fatal(err)
}
t.Logf("rows after healthy createTask: batches=%d jobs=%d", batches, jobs)
if batches != 1 || jobs != 2 {
t.Fatalf("a healthy grant produced batches=%d jobs=%d, want 1 and 2", batches, jobs)
}
g.mu.Lock()
defer g.mu.Unlock()
if g.refreshes != 1 {
t.Fatalf("token refreshes = %d, want exactly the one the probe needed", g.refreshes)
}
if len(g.mintCalls) != 1 || g.mintCalls[0] != "org.tangled.repo.createRepo aud="+serviceDid {
t.Fatalf("probe mint calls = %v, want one createRepo mint for the migrator's own did", g.mintCalls)
}
}