mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
feat(api-key): show last used IP
This commit is contained in:
@@ -10,6 +10,7 @@ import (
|
||||
|
||||
func TestAPIKeyFromService_MapsLastUsedAt(t *testing.T) {
|
||||
lastUsed := time.Now().UTC().Truncate(time.Second)
|
||||
lastUsedIP := "203.0.113.10"
|
||||
src := &service.APIKey{
|
||||
ID: 1,
|
||||
UserID: 2,
|
||||
@@ -17,6 +18,7 @@ func TestAPIKeyFromService_MapsLastUsedAt(t *testing.T) {
|
||||
Name: "Mapper",
|
||||
Status: service.StatusActive,
|
||||
LastUsedAt: &lastUsed,
|
||||
LastUsedIP: &lastUsedIP,
|
||||
CurrentConcurrency: 3,
|
||||
}
|
||||
|
||||
@@ -24,6 +26,8 @@ func TestAPIKeyFromService_MapsLastUsedAt(t *testing.T) {
|
||||
require.NotNil(t, out)
|
||||
require.NotNil(t, out.LastUsedAt)
|
||||
require.WithinDuration(t, lastUsed, *out.LastUsedAt, time.Second)
|
||||
require.NotNil(t, out.LastUsedIP)
|
||||
require.Equal(t, lastUsedIP, *out.LastUsedIP)
|
||||
require.Equal(t, 3, out.CurrentConcurrency)
|
||||
}
|
||||
|
||||
@@ -39,4 +43,5 @@ func TestAPIKeyFromService_MapsNilLastUsedAt(t *testing.T) {
|
||||
out := APIKeyFromService(src)
|
||||
require.NotNil(t, out)
|
||||
require.Nil(t, out.LastUsedAt)
|
||||
require.Nil(t, out.LastUsedIP)
|
||||
}
|
||||
|
||||
@@ -89,6 +89,7 @@ func APIKeyFromService(k *service.APIKey) *APIKey {
|
||||
IPWhitelist: k.IPWhitelist,
|
||||
IPBlacklist: k.IPBlacklist,
|
||||
LastUsedAt: k.LastUsedAt,
|
||||
LastUsedIP: k.LastUsedIP,
|
||||
Quota: k.Quota,
|
||||
QuotaUsed: k.QuotaUsed,
|
||||
ExpiresAt: k.ExpiresAt,
|
||||
|
||||
@@ -59,6 +59,7 @@ type APIKey struct {
|
||||
IPWhitelist []string `json:"ip_whitelist"`
|
||||
IPBlacklist []string `json:"ip_blacklist"`
|
||||
LastUsedAt *time.Time `json:"last_used_at"`
|
||||
LastUsedIP *string `json:"last_used_ip"`
|
||||
Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
|
||||
QuotaUsed float64 `json:"quota_used"` // Used quota amount in USD
|
||||
ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = never expires)
|
||||
|
||||
@@ -14,9 +14,11 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/ent/schema/mixins"
|
||||
"github.com/Wei-Shaw/sub2api/ent/user"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/lib/pq"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
|
||||
"entgo.io/ent/dialect"
|
||||
entsql "entgo.io/ent/dialect/sql"
|
||||
)
|
||||
|
||||
@@ -431,10 +433,96 @@ func (r *apiKeyRepository) ListByUserID(ctx context.Context, userID int64, param
|
||||
for i := range keys {
|
||||
outKeys = append(outKeys, *apiKeyEntityToService(keys[i]))
|
||||
}
|
||||
if err := r.attachLastUsedIPs(ctx, outKeys); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
return outKeys, paginationResultFromTotal(int64(total), params), nil
|
||||
}
|
||||
|
||||
func (r *apiKeyRepository) attachLastUsedIPs(ctx context.Context, keys []service.APIKey) error {
|
||||
if len(keys) == 0 || r.sql == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
apiKeyIDs := make([]int64, 0, len(keys))
|
||||
for i := range keys {
|
||||
apiKeyIDs = append(apiKeyIDs, keys[i].ID)
|
||||
}
|
||||
|
||||
lastUsedIPs, err := r.latestUsageLogIPs(ctx, apiKeyIDs)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for i := range keys {
|
||||
if ip, ok := lastUsedIPs[keys[i].ID]; ok {
|
||||
keys[i].LastUsedIP = &ip
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *apiKeyRepository) latestUsageLogIPs(ctx context.Context, apiKeyIDs []int64) (map[int64]string, error) {
|
||||
if len(apiKeyIDs) == 0 || r.sql == nil {
|
||||
return map[int64]string{}, nil
|
||||
}
|
||||
|
||||
query, args := latestUsageLogIPsQuery(apiKeyIDs, r.client.Driver().Dialect())
|
||||
rows, err := r.sql.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := make(map[int64]string, len(apiKeyIDs))
|
||||
for rows.Next() {
|
||||
var apiKeyID int64
|
||||
var ipAddress string
|
||||
if err := rows.Scan(&apiKeyID, &ipAddress); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[apiKeyID] = ipAddress
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func latestUsageLogIPsQuery(apiKeyIDs []int64, dialectName string) (string, []any) {
|
||||
if dialectName == dialect.Postgres {
|
||||
return `
|
||||
SELECT api_key_id, ip_address
|
||||
FROM (
|
||||
SELECT api_key_id, ip_address,
|
||||
ROW_NUMBER() OVER (PARTITION BY api_key_id ORDER BY created_at DESC, id DESC) AS rn
|
||||
FROM usage_logs
|
||||
WHERE api_key_id = ANY($1::bigint[])
|
||||
AND ip_address IS NOT NULL
|
||||
AND ip_address <> ''
|
||||
) ranked
|
||||
WHERE rn = 1`, []any{pq.Array(apiKeyIDs)}
|
||||
}
|
||||
|
||||
placeholders := make([]string, len(apiKeyIDs))
|
||||
args := make([]any, len(apiKeyIDs))
|
||||
for i, id := range apiKeyIDs {
|
||||
placeholders[i] = "?"
|
||||
args[i] = id
|
||||
}
|
||||
return fmt.Sprintf(`
|
||||
SELECT api_key_id, ip_address
|
||||
FROM (
|
||||
SELECT api_key_id, ip_address,
|
||||
ROW_NUMBER() OVER (PARTITION BY api_key_id ORDER BY created_at DESC, id DESC) AS rn
|
||||
FROM usage_logs
|
||||
WHERE api_key_id IN (%s)
|
||||
AND ip_address IS NOT NULL
|
||||
AND ip_address <> ''
|
||||
) ranked
|
||||
WHERE rn = 1`, strings.Join(placeholders, ", ")), args
|
||||
}
|
||||
|
||||
func (r *apiKeyRepository) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
|
||||
if len(apiKeyIDs) == 0 {
|
||||
return []int64{}, nil
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/ent/enttest"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
@@ -30,7 +31,7 @@ func newAPIKeyRepoSQLite(t *testing.T) (*apiKeyRepository, *dbent.Client) {
|
||||
client := enttest.NewClient(t, enttest.WithOptions(dbent.Driver(drv)))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
|
||||
return &apiKeyRepository{client: client}, client
|
||||
return &apiKeyRepository{client: client, sql: db}, client
|
||||
}
|
||||
|
||||
func mustCreateAPIKeyRepoUser(t *testing.T, ctx context.Context, client *dbent.Client, email string) *service.User {
|
||||
@@ -45,6 +46,85 @@ func mustCreateAPIKeyRepoUser(t *testing.T, ctx context.Context, client *dbent.C
|
||||
return userEntityToService(u)
|
||||
}
|
||||
|
||||
func mustCreateAPIKeyRepoAccount(t *testing.T, ctx context.Context, client *dbent.Client, name string) int64 {
|
||||
t.Helper()
|
||||
a, err := client.Account.Create().
|
||||
SetName(name).
|
||||
SetPlatform(service.PlatformOpenAI).
|
||||
SetType(service.AccountTypeAPIKey).
|
||||
SetStatus(service.StatusActive).
|
||||
SetCredentials(map[string]any{"api_key": "sk-test"}).
|
||||
Save(ctx)
|
||||
require.NoError(t, err)
|
||||
return a.ID
|
||||
}
|
||||
|
||||
func mustCreateAPIKeyRepoUsageLog(t *testing.T, ctx context.Context, client *dbent.Client, userID, apiKeyID, accountID int64, requestID string, createdAt time.Time, ipAddress *string) {
|
||||
t.Helper()
|
||||
builder := client.UsageLog.Create().
|
||||
SetUserID(userID).
|
||||
SetAPIKeyID(apiKeyID).
|
||||
SetAccountID(accountID).
|
||||
SetRequestID(requestID).
|
||||
SetModel("gpt-5").
|
||||
SetCreatedAt(createdAt)
|
||||
if ipAddress != nil {
|
||||
builder.SetIPAddress(*ipAddress)
|
||||
}
|
||||
_, err := builder.Save(ctx)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestAPIKeyRepositoryListByUserIDAttachesLastUsedIP(t *testing.T) {
|
||||
repo, client := newAPIKeyRepoSQLite(t)
|
||||
ctx := context.Background()
|
||||
user := mustCreateAPIKeyRepoUser(t, ctx, client, "list-last-used-ip@test.com")
|
||||
accountID := mustCreateAPIKeyRepoAccount(t, ctx, client, "acc-list-last-used-ip")
|
||||
|
||||
withLogs := &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-list-last-used-ip-logs",
|
||||
Name: "With Logs",
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
emptyOnly := &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-list-last-used-ip-empty",
|
||||
Name: "Empty Only",
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
noLogs := &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-list-last-used-ip-none",
|
||||
Name: "No Logs",
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, withLogs))
|
||||
require.NoError(t, repo.Create(ctx, emptyOnly))
|
||||
require.NoError(t, repo.Create(ctx, noLogs))
|
||||
|
||||
olderIP := "198.51.100.10"
|
||||
newerEmptyIP := ""
|
||||
newestIP := "203.0.113.20"
|
||||
base := time.Now().UTC().Add(-3 * time.Hour).Truncate(time.Second)
|
||||
mustCreateAPIKeyRepoUsageLog(t, ctx, client, user.ID, withLogs.ID, accountID, "req-last-ip-older", base, &olderIP)
|
||||
mustCreateAPIKeyRepoUsageLog(t, ctx, client, user.ID, withLogs.ID, accountID, "req-last-ip-empty", base.Add(time.Hour), &newerEmptyIP)
|
||||
mustCreateAPIKeyRepoUsageLog(t, ctx, client, user.ID, withLogs.ID, accountID, "req-last-ip-newest", base.Add(2*time.Hour), &newestIP)
|
||||
mustCreateAPIKeyRepoUsageLog(t, ctx, client, user.ID, emptyOnly.ID, accountID, "req-empty-ip", base.Add(3*time.Hour), &newerEmptyIP)
|
||||
|
||||
keys, _, err := repo.ListByUserID(ctx, user.ID, pagination.PaginationParams{Page: 1, PageSize: 10}, service.APIKeyListFilters{})
|
||||
require.NoError(t, err)
|
||||
|
||||
byID := make(map[int64]service.APIKey, len(keys))
|
||||
for _, key := range keys {
|
||||
byID[key.ID] = key
|
||||
}
|
||||
require.NotNil(t, byID[withLogs.ID].LastUsedIP)
|
||||
require.Equal(t, newestIP, *byID[withLogs.ID].LastUsedIP)
|
||||
require.Nil(t, byID[emptyOnly.ID].LastUsedIP)
|
||||
require.Nil(t, byID[noLogs.ID].LastUsedIP)
|
||||
}
|
||||
|
||||
func TestAPIKeyRepository_CreateWithLastUsedAt(t *testing.T) {
|
||||
repo, client := newAPIKeyRepoSQLite(t)
|
||||
ctx := context.Background()
|
||||
|
||||
@@ -234,6 +234,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"ip_whitelist": null,
|
||||
"ip_blacklist": null,
|
||||
"last_used_at": null,
|
||||
"last_used_ip": null,
|
||||
"current_concurrency": 0,
|
||||
"quota": 0,
|
||||
"quota_used": 0,
|
||||
@@ -284,6 +285,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"ip_whitelist": null,
|
||||
"ip_blacklist": null,
|
||||
"last_used_at": null,
|
||||
"last_used_ip": null,
|
||||
"current_concurrency": 0,
|
||||
"quota": 0,
|
||||
"quota_used": 0,
|
||||
|
||||
@@ -40,6 +40,7 @@ type APIKey struct {
|
||||
CompiledIPWhitelist *ip.CompiledIPRules `json:"-"`
|
||||
CompiledIPBlacklist *ip.CompiledIPRules `json:"-"`
|
||||
LastUsedAt *time.Time
|
||||
LastUsedIP *string
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
User *User
|
||||
|
||||
Reference in New Issue
Block a user