Skip to content

Commit a140235

Browse files
committed
Add local mail API rate limiting
Block configured mail OpenAPI bursts before token resolution, sharing a workspace-local rolling window across CLI processes. Preserve local API errors across raw and generated command paths so rate_limit remains visible to callers. sprint: S2
1 parent 020aeb8 commit a140235

18 files changed

Lines changed: 1210 additions & 6 deletions

cmd/api/api.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -276,7 +276,7 @@ func apiPaginate(ctx context.Context, ac *client.APIClient, request client.RawAp
276276
pf.FormatPage(items)
277277
}, pagOpts)
278278
if err != nil {
279-
return output.MarkRaw(output.ErrNetwork("API call failed: %v", err))
279+
return output.MarkRaw(client.WrapDoAPIError(err))
280280
}
281281
if apiErr := client.CheckLarkResponse(result); apiErr != nil {
282282
output.FormatValue(out, result, output.FormatJSON)
@@ -290,7 +290,7 @@ func apiPaginate(ctx context.Context, ac *client.APIClient, request client.RawAp
290290
default:
291291
result, err := ac.PaginateAll(ctx, request, pagOpts)
292292
if err != nil {
293-
return output.MarkRaw(output.ErrNetwork("API call failed: %v", err))
293+
return output.MarkRaw(client.WrapDoAPIError(err))
294294
}
295295
if apiErr := client.CheckLarkResponse(result); apiErr != nil {
296296
output.FormatValue(out, result, output.FormatJSON)

cmd/api/api_test.go

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,18 @@
44
package api
55

66
import (
7+
"context"
78
"errors"
89
"os"
910
"sort"
1011
"strings"
1112
"testing"
13+
"time"
1214

1315
"github.com/larksuite/cli/internal/cmdutil"
1416
"github.com/larksuite/cli/internal/core"
1517
"github.com/larksuite/cli/internal/httpmock"
18+
"github.com/larksuite/cli/internal/mailratelimit"
1619
"github.com/larksuite/cli/internal/output"
1720
"github.com/spf13/cobra"
1821
)
@@ -66,6 +69,46 @@ func TestApiCmd_DryRun(t *testing.T) {
6669
}
6770
}
6871

72+
func TestApiCmd_MailRateLimitErrorIsRawRateLimit(t *testing.T) {
73+
installLimitedMailWindow(t)
74+
f, _, _, _ := cmdutil.TestFactory(t, &core.CliConfig{
75+
AppID: "test-app", AppSecret: "test-secret", Brand: core.BrandFeishu,
76+
})
77+
78+
cmd := NewCmdApi(f, nil)
79+
cmd.SetArgs([]string{"GET", "/open-apis/mail/v1/user_mailboxes/me/messages/msg_1", "--as", "bot"})
80+
err := cmd.Execute()
81+
if err == nil {
82+
t.Fatal("expected rate limit error")
83+
}
84+
var exitErr *output.ExitError
85+
if !errors.As(err, &exitErr) {
86+
t.Fatalf("expected ExitError, got %T", err)
87+
}
88+
if exitErr.Detail == nil || exitErr.Detail.Type != "rate_limit" || exitErr.Detail.Code != output.LarkErrRateLimit {
89+
t.Fatalf("unexpected detail: %#v", exitErr.Detail)
90+
}
91+
if !exitErr.Raw {
92+
t.Fatal("expected raw API error to be marked raw")
93+
}
94+
}
95+
96+
func TestApiCmd_MailDryRunBypassesRateLimit(t *testing.T) {
97+
installLimitedMailWindow(t)
98+
f, stdout, _, _ := cmdutil.TestFactory(t, &core.CliConfig{
99+
AppID: "test-app", AppSecret: "test-secret", Brand: core.BrandFeishu,
100+
})
101+
102+
cmd := NewCmdApi(f, nil)
103+
cmd.SetArgs([]string{"GET", "/open-apis/mail/v1/user_mailboxes/me/messages/msg_1", "--as", "bot", "--dry-run"})
104+
if err := cmd.Execute(); err != nil {
105+
t.Fatalf("dry-run err = %v", err)
106+
}
107+
if !strings.Contains(stdout.String(), "Dry Run") {
108+
t.Fatalf("expected dry-run output, got %s", stdout.String())
109+
}
110+
}
111+
69112
func TestApiCmd_BotMode(t *testing.T) {
70113
f, stdout, _, reg := cmdutil.TestFactory(t, &core.CliConfig{
71114
AppID: "test-app", AppSecret: "test-secret", Brand: core.BrandFeishu,
@@ -88,6 +131,29 @@ func TestApiCmd_BotMode(t *testing.T) {
88131
}
89132
}
90133

134+
func installLimitedMailWindow(t *testing.T) {
135+
t.Helper()
136+
rule := mailratelimit.Rule{
137+
Method: "GET",
138+
CanonicalPath: "/open-apis/mail/v1/user_mailboxes/:user_mailbox_id/messages/:message_id",
139+
Window: 2 * time.Second,
140+
Limit: 1,
141+
Scope: mailratelimit.ScopeApp,
142+
}
143+
now := time.Unix(100, 0)
144+
limiter := mailratelimit.NewLimiterForDir(t.TempDir(), []mailratelimit.Rule{rule}, func() time.Time { return now })
145+
restore := mailratelimit.SetDefaultLimiterForTest(limiter)
146+
t.Cleanup(restore)
147+
if err := mailratelimit.Check(context.Background(), mailratelimit.Request{
148+
Brand: core.BrandFeishu,
149+
AppID: "test-app",
150+
Method: "GET",
151+
Path: "/open-apis/mail/v1/user_mailboxes/me/messages/msg_1",
152+
}); err != nil {
153+
t.Fatalf("prefill limiter: %v", err)
154+
}
155+
}
156+
91157
func TestApiCmd_MissingArgs(t *testing.T) {
92158
f, _, _, _ := cmdutil.TestFactory(t, &core.CliConfig{
93159
AppID: "test-app", AppSecret: "test-secret", Brand: core.BrandFeishu,

cmd/service/service.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -280,7 +280,7 @@ func serviceMethodRun(opts *ServiceMethodOptions) error {
280280

281281
resp, err := ac.DoAPI(opts.Ctx, request)
282282
if err != nil {
283-
return output.ErrNetwork("API call failed: %s", err)
283+
return client.WrapDoAPIError(err)
284284
}
285285
return client.HandleResponse(resp, client.ResponseOptions{
286286
OutputPath: opts.Output,
@@ -516,7 +516,7 @@ func servicePaginate(ctx context.Context, ac *client.APIClient, request client.R
516516
pf.FormatPage(items)
517517
}, pagOpts)
518518
if err != nil {
519-
return output.ErrNetwork("API call failed: %s", err)
519+
return client.WrapDoAPIError(err)
520520
}
521521
if apiErr := checkErr(result); apiErr != nil {
522522
return apiErr
@@ -529,7 +529,7 @@ func servicePaginate(ctx context.Context, ac *client.APIClient, request client.R
529529
default:
530530
result, err := ac.PaginateAll(ctx, request, pagOpts)
531531
if err != nil {
532-
return output.ErrNetwork("API call failed: %s", err)
532+
return client.WrapDoAPIError(err)
533533
}
534534
if apiErr := checkErr(result); apiErr != nil {
535535
return apiErr

cmd/service/service_test.go

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,17 @@
44
package service
55

66
import (
7+
"context"
8+
"errors"
79
"os"
810
"strings"
911
"testing"
12+
"time"
1013

1114
"github.com/larksuite/cli/internal/cmdutil"
1215
"github.com/larksuite/cli/internal/core"
1316
"github.com/larksuite/cli/internal/httpmock"
17+
"github.com/larksuite/cli/internal/mailratelimit"
1418
"github.com/larksuite/cli/internal/output"
1519
"github.com/spf13/cobra"
1620
)
@@ -45,6 +49,28 @@ func driveMethod(httpMethod string, params map[string]interface{}) map[string]in
4549
return m
4650
}
4751

52+
func mailMessageSpec() map[string]interface{} {
53+
return map[string]interface{}{
54+
"name": "mail",
55+
"servicePath": "/open-apis/mail/v1",
56+
}
57+
}
58+
59+
func mailMessageMethod() map[string]interface{} {
60+
return map[string]interface{}{
61+
"path": "user_mailboxes/{user_mailbox_id}/messages/{message_id}",
62+
"httpMethod": "GET",
63+
"parameters": map[string]interface{}{
64+
"user_mailbox_id": map[string]interface{}{
65+
"type": "string", "location": "path", "required": true,
66+
},
67+
"message_id": map[string]interface{}{
68+
"type": "string", "location": "path", "required": true,
69+
},
70+
},
71+
}
72+
}
73+
4874
// ── registerService ──
4975

5076
func TestRegisterService(t *testing.T) {
@@ -192,6 +218,72 @@ func TestNewCmdServiceMethod_RunFCallback(t *testing.T) {
192218
}
193219
}
194220

221+
func TestServiceMethod_MailRateLimitErrorIsNotNetwork(t *testing.T) {
222+
installLimitedMailWindow(t)
223+
f, _, _, _ := cmdutil.TestFactory(t, testConfig)
224+
225+
cmd := NewCmdServiceMethod(f, mailMessageSpec(), mailMessageMethod(), "get", "messages", nil)
226+
cmd.SetArgs([]string{
227+
"--as", "bot",
228+
"--params", `{"user_mailbox_id":"me","message_id":"msg_1"}`,
229+
})
230+
err := cmd.Execute()
231+
if err == nil {
232+
t.Fatal("expected rate limit error")
233+
}
234+
var exitErr *output.ExitError
235+
if !errors.As(err, &exitErr) {
236+
t.Fatalf("expected ExitError, got %T", err)
237+
}
238+
if exitErr.Detail == nil || exitErr.Detail.Type != "rate_limit" || exitErr.Detail.Code != output.LarkErrRateLimit {
239+
t.Fatalf("unexpected detail: %#v", exitErr.Detail)
240+
}
241+
if exitErr.Detail.Type == "network" {
242+
t.Fatal("rate limit was wrapped as network")
243+
}
244+
}
245+
246+
func TestServiceMethod_MailDryRunBypassesRateLimit(t *testing.T) {
247+
installLimitedMailWindow(t)
248+
f, stdout, _, _ := cmdutil.TestFactory(t, testConfig)
249+
250+
cmd := NewCmdServiceMethod(f, mailMessageSpec(), mailMessageMethod(), "get", "messages", nil)
251+
cmd.SetArgs([]string{
252+
"--as", "bot",
253+
"--params", `{"user_mailbox_id":"me","message_id":"msg_1"}`,
254+
"--dry-run",
255+
})
256+
if err := cmd.Execute(); err != nil {
257+
t.Fatalf("dry-run err = %v", err)
258+
}
259+
if !strings.Contains(stdout.String(), "Dry Run") {
260+
t.Fatalf("expected dry-run output, got %s", stdout.String())
261+
}
262+
}
263+
264+
func installLimitedMailWindow(t *testing.T) {
265+
t.Helper()
266+
rule := mailratelimit.Rule{
267+
Method: "GET",
268+
CanonicalPath: "/open-apis/mail/v1/user_mailboxes/:user_mailbox_id/messages/:message_id",
269+
Window: 2 * time.Second,
270+
Limit: 1,
271+
Scope: mailratelimit.ScopeApp,
272+
}
273+
now := time.Unix(100, 0)
274+
limiter := mailratelimit.NewLimiterForDir(t.TempDir(), []mailratelimit.Rule{rule}, func() time.Time { return now })
275+
restore := mailratelimit.SetDefaultLimiterForTest(limiter)
276+
t.Cleanup(restore)
277+
if err := mailratelimit.Check(context.Background(), mailratelimit.Request{
278+
Brand: core.BrandFeishu,
279+
AppID: "test-app",
280+
Method: "GET",
281+
Path: "/open-apis/mail/v1/user_mailboxes/me/messages/msg_1",
282+
}); err != nil {
283+
t.Fatalf("prefill limiter: %v", err)
284+
}
285+
}
286+
195287
// ── dry-run / buildServiceRequest ──
196288

197289
func TestServiceMethod_DryRun_PathParam(t *testing.T) {

internal/client/api_errors.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,10 @@ func WrapDoAPIError(err error) error {
2323
if err == nil {
2424
return nil
2525
}
26+
var exitErr *output.ExitError
27+
if errors.As(err, &exitErr) {
28+
return err
29+
}
2630
if isJSONDecodeError(err, false) {
2731
return output.ErrWithHint(output.ExitAPI, "api_error",
2832
fmt.Sprintf("API returned an invalid JSON response: %v", err), rawAPIJSONHint)

internal/client/api_errors_test.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,18 @@ func TestWrapDoAPIError_SyntaxErrorIsAPIDiagnostic(t *testing.T) {
4949
}
5050
}
5151

52+
func TestWrapDoAPIError_PreservesExitError(t *testing.T) {
53+
original := output.ErrAPI(output.LarkErrRateLimit, "rate limited", map[string]any{"retry_after_ms": 100})
54+
err := WrapDoAPIError(original)
55+
if err != original {
56+
t.Fatalf("WrapDoAPIError returned %p, want original %p", err, original)
57+
}
58+
var exitErr *output.ExitError
59+
if !errors.As(err, &exitErr) || exitErr.Detail == nil || exitErr.Detail.Type != "rate_limit" {
60+
t.Fatalf("err = %v, want preserved rate_limit ExitError", err)
61+
}
62+
}
63+
5264
func TestWrapJSONResponseParseError_UnexpectedEOFIsAPIDiagnostic(t *testing.T) {
5365
err := WrapJSONResponseParseError(io.ErrUnexpectedEOF, []byte("{"))
5466
if err == nil {

internal/client/client.go

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020

2121
"github.com/larksuite/cli/internal/core"
2222
"github.com/larksuite/cli/internal/credential"
23+
"github.com/larksuite/cli/internal/mailratelimit"
2324
"github.com/larksuite/cli/internal/output"
2425
"github.com/larksuite/cli/internal/util"
2526
)
@@ -94,6 +95,10 @@ func (c *APIClient) buildApiReq(request RawApiRequest) (*larkcore.ApiReq, []lark
9495
func (c *APIClient) DoSDKRequest(ctx context.Context, req *larkcore.ApiReq, as core.Identity, extraOpts ...larkcore.RequestOptionFunc) (*larkcore.ApiResp, error) {
9596
var opts []larkcore.RequestOptionFunc
9697

98+
if err := c.checkMailRateLimit(ctx, req); err != nil {
99+
return nil, err
100+
}
101+
97102
token, err := c.resolveAccessToken(ctx, as)
98103
if err != nil {
99104
return nil, err
@@ -120,6 +125,10 @@ func (c *APIClient) DoSDKRequest(ctx context.Context, req *larkcore.ApiReq, as c
120125
func (c *APIClient) DoStream(ctx context.Context, req *larkcore.ApiReq, as core.Identity, opts ...Option) (*http.Response, error) {
121126
cfg := buildConfig(opts)
122127

128+
if err := c.checkMailRateLimit(ctx, req); err != nil {
129+
return nil, err
130+
}
131+
123132
// Resolve auth
124133
token, err := c.resolveAccessToken(ctx, as)
125134
if err != nil {
@@ -190,6 +199,21 @@ func (c *APIClient) DoStream(ctx context.Context, req *larkcore.ApiReq, as core.
190199
return resp, nil
191200
}
192201

202+
func (c *APIClient) checkMailRateLimit(ctx context.Context, req *larkcore.ApiReq) error {
203+
if req == nil {
204+
return output.ErrValidation("api request is required")
205+
}
206+
limitReq := mailratelimit.Request{
207+
Method: req.HttpMethod,
208+
Path: req.ApiPath,
209+
}
210+
if c != nil && c.Config != nil {
211+
limitReq.Brand = c.Config.Brand
212+
limitReq.AppID = c.Config.AppID
213+
}
214+
return mailratelimit.Check(ctx, limitReq)
215+
}
216+
193217
type cancelOnCloseBody struct {
194218
io.ReadCloser
195219
cancel context.CancelFunc

0 commit comments

Comments
 (0)