feat(api-key): show last used IP

This commit is contained in:
wucm667
2026-07-08 15:26:54 +08:00
parent 6f43986c37
commit e0d149d511
12 changed files with 245 additions and 6 deletions
@@ -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)
}
+1
View File
@@ -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,
+1
View File
@@ -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,
+1
View File
@@ -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