4
0

middleware_test.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662
  1. package webapi
  2. import (
  3. "context"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "log/slog"
  8. "net/http"
  9. "net/http/httptest"
  10. "strings"
  11. "testing"
  12. "time"
  13. "github.com/stretchr/testify/assert"
  14. "github.com/stretchr/testify/require"
  15. "github.com/mk6i/open-oscar-server/state"
  16. "github.com/mk6i/open-oscar-server/wire"
  17. )
  18. // stubValidator records how many times a key was looked up so the tests can
  19. // assert that CORSMiddleware and the auth layer share a single lookup.
  20. type stubValidator struct {
  21. key *state.WebAPIKey
  22. lookup int
  23. }
  24. func (s *stubValidator) GetAPIKeyByDevKey(_ context.Context, devKey string) (*state.WebAPIKey, error) {
  25. s.lookup++
  26. if s.key == nil || s.key.DevKey != devKey {
  27. return nil, nil
  28. }
  29. return s.key, nil
  30. }
  31. func newTestMiddleware(v APIKeyValidator) *AuthMiddleware {
  32. return NewAuthMiddleware(v, slog.New(slog.NewTextHandler(io.Discard, nil)))
  33. }
  34. // The Web AIM client permanently downgrades to JSONP when a cross-origin
  35. // response arrives without Access-Control-Allow-Origin, so every response the
  36. // auth layer rejects must still carry CORS headers. That only holds while
  37. // CORSMiddleware wraps the auth middleware.
  38. func TestCORSMiddleware_HeadersOnAuthRejection(t *testing.T) {
  39. // The auth layer reports failures in the response envelope rather than the
  40. // HTTP status, so the envelope's statusCode is what identifies a rejection.
  41. tests := []struct {
  42. name string
  43. query string
  44. validator *stubValidator
  45. wantEnvelope string
  46. }{
  47. {
  48. name: "missing credentials",
  49. query: "?f=json",
  50. validator: &stubValidator{},
  51. wantEnvelope: `"statusCode":400`,
  52. },
  53. {
  54. name: "unknown api key",
  55. query: "?k=nosuchkey",
  56. validator: &stubValidator{},
  57. wantEnvelope: `"statusCode":403`,
  58. },
  59. }
  60. for _, tt := range tests {
  61. t.Run(tt.name, func(t *testing.T) {
  62. m := newTestMiddleware(tt.validator)
  63. var reachedHandler bool
  64. h := m.CORSMiddleware(m.AuthenticateFlexible(
  65. http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
  66. reachedHandler = true
  67. w.WriteHeader(http.StatusOK)
  68. })))
  69. r := httptest.NewRequest(http.MethodGet, "/im/sendIM"+tt.query, nil)
  70. r.Header.Set("Origin", "http://localhost:8000")
  71. w := httptest.NewRecorder()
  72. h.ServeHTTP(w, r)
  73. assert.False(t, reachedHandler, "auth layer should have rejected the request")
  74. assert.Contains(t, w.Body.String(), tt.wantEnvelope)
  75. assert.Equal(t, "http://localhost:8000", w.Header().Get("Access-Control-Allow-Origin"))
  76. assert.Equal(t, "Origin", w.Header().Get("Vary"))
  77. })
  78. }
  79. }
  80. // A 404 for an endpoint this server does not implement (/service/getAttributes,
  81. // /metrics/sendIM) must reach the client as a 404 rather than as a blocked
  82. // response.
  83. func TestCORSMiddleware_HeadersOn404(t *testing.T) {
  84. m := newTestMiddleware(&stubValidator{})
  85. h := m.CORSMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
  86. w.WriteHeader(http.StatusNotFound)
  87. }))
  88. r := httptest.NewRequest(http.MethodGet, "/service/getAttributes?f=json&aimsid=abc", nil)
  89. r.Header.Set("Origin", "http://localhost:8000")
  90. w := httptest.NewRecorder()
  91. h.ServeHTTP(w, r)
  92. assert.Equal(t, http.StatusNotFound, w.Code)
  93. assert.Equal(t, "http://localhost:8000", w.Header().Get("Access-Control-Allow-Origin"))
  94. }
  95. func TestCORSMiddleware_PreflightShortCircuits(t *testing.T) {
  96. m := newTestMiddleware(&stubValidator{})
  97. var reachedNext bool
  98. h := m.CORSMiddleware(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
  99. reachedNext = true
  100. }))
  101. r := httptest.NewRequest(http.MethodOptions, "/im/sendIM", nil)
  102. r.Header.Set("Origin", "http://localhost:8000")
  103. w := httptest.NewRecorder()
  104. h.ServeHTTP(w, r)
  105. assert.Equal(t, http.StatusNoContent, w.Code)
  106. assert.False(t, reachedNext, "preflight must not reach the wrapped handler")
  107. assert.Equal(t, "http://localhost:8000", w.Header().Get("Access-Control-Allow-Origin"))
  108. assert.Contains(t, w.Header().Get("Access-Control-Allow-Methods"), "POST")
  109. }
  110. // CORSMiddleware runs ahead of authentication and so resolves the API key
  111. // itself; the auth layer behind it must reuse that lookup rather than repeat it.
  112. func TestCORSMiddleware_SharesKeyLookupWithAuth(t *testing.T) {
  113. v := &stubValidator{key: &state.WebAPIKey{
  114. DevID: "dev1",
  115. DevKey: "goodkey",
  116. IsActive: true,
  117. RateLimit: 100,
  118. }}
  119. m := newTestMiddleware(v)
  120. var served bool
  121. h := m.CORSMiddleware(m.Authenticate(
  122. http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  123. served = true
  124. key, ok := r.Context().Value(ContextKeyAPIKey).(*state.WebAPIKey)
  125. require.True(t, ok, "handler should see the validated key")
  126. assert.Equal(t, "dev1", key.DevID)
  127. w.WriteHeader(http.StatusOK)
  128. })))
  129. r := httptest.NewRequest(http.MethodGet, "/aim/startOSCARSession?k=goodkey", nil)
  130. r.Header.Set("Origin", "http://localhost:8000")
  131. w := httptest.NewRecorder()
  132. h.ServeHTTP(w, r)
  133. assert.True(t, served)
  134. assert.Equal(t, http.StatusOK, w.Code)
  135. assert.Equal(t, 1, v.lookup, "key should be resolved once per request, not once per middleware")
  136. }
  137. // Per-key origin allowlists must keep working now that the origin decision is
  138. // made before authentication.
  139. func TestCORSMiddleware_PerKeyOriginAllowlist(t *testing.T) {
  140. v := &stubValidator{key: &state.WebAPIKey{
  141. DevID: "dev1",
  142. DevKey: "goodkey",
  143. IsActive: true,
  144. RateLimit: 100,
  145. AllowedOrigins: []string{"http://allowed.example"},
  146. }}
  147. m := newTestMiddleware(v)
  148. h := m.CORSMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
  149. w.WriteHeader(http.StatusOK)
  150. }))
  151. t.Run("allowed origin", func(t *testing.T) {
  152. r := httptest.NewRequest(http.MethodGet, "/im/sendIM?k=goodkey", nil)
  153. r.Header.Set("Origin", "http://allowed.example")
  154. w := httptest.NewRecorder()
  155. h.ServeHTTP(w, r)
  156. assert.Equal(t, "http://allowed.example", w.Header().Get("Access-Control-Allow-Origin"))
  157. })
  158. t.Run("disallowed origin", func(t *testing.T) {
  159. r := httptest.NewRequest(http.MethodGet, "/im/sendIM?k=goodkey", nil)
  160. r.Header.Set("Origin", "http://evil.example")
  161. w := httptest.NewRecorder()
  162. h.ServeHTTP(w, r)
  163. assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"))
  164. })
  165. }
  166. // A POST body must survive the middleware chain: CORSMiddleware reads the API
  167. // key from the query string only, so it never parses the form.
  168. func TestCORSMiddleware_DoesNotConsumePOSTBody(t *testing.T) {
  169. m := newTestMiddleware(&stubValidator{})
  170. h := m.CORSMiddleware(m.AuthenticateFlexible(
  171. http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  172. body, err := io.ReadAll(r.Body)
  173. require.NoError(t, err)
  174. assert.Equal(t, "message=hello", string(body))
  175. w.WriteHeader(http.StatusOK)
  176. })))
  177. r := httptest.NewRequest(http.MethodPost, "/im/sendIM?aimsid=abc", strings.NewReader("message=hello"))
  178. r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
  179. r.Header.Set("Origin", "http://localhost:8000")
  180. w := httptest.NewRecorder()
  181. h.ServeHTTP(w, r)
  182. assert.Equal(t, http.StatusOK, w.Code)
  183. }
  184. // Some clients POST the whole parameter set in the body, so an auth layer reading
  185. // only the query string sees no credential. Reading the body here does not cost the
  186. // handler its parameters: ParseForm caches onto the request.
  187. func TestAuthenticateFlexible_ReadsCredentialsFromPOSTBody(t *testing.T) {
  188. tests := []struct {
  189. name string
  190. body string
  191. contentType string
  192. }{
  193. {
  194. name: "declared form body",
  195. body: "aimsid=abc&message=hello",
  196. contentType: "application/x-www-form-urlencoded",
  197. },
  198. {
  199. // No Content-Type announced; the request is form data all the same.
  200. name: "untyped form body",
  201. body: "aimsid=abc&message=hello",
  202. },
  203. }
  204. for _, tt := range tests {
  205. t.Run(tt.name, func(t *testing.T) {
  206. m := newTestMiddleware(&stubValidator{})
  207. reached := false
  208. h := m.CORSMiddleware(m.AuthenticateFlexible(
  209. http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  210. reached = true
  211. assert.Equal(t, "hello", param(r, "message"))
  212. w.WriteHeader(http.StatusOK)
  213. })))
  214. r := httptest.NewRequest(http.MethodPost, "/im/sendIM", strings.NewReader(tt.body))
  215. if tt.contentType != "" {
  216. r.Header.Set("Content-Type", tt.contentType)
  217. }
  218. w := httptest.NewRecorder()
  219. h.ServeHTTP(w, r)
  220. assert.True(t, reached, "request was rejected before reaching the handler")
  221. assert.Equal(t, http.StatusOK, w.Code)
  222. })
  223. }
  224. }
  225. // The auth layer's own rejections must be JSONP-wrapped too, otherwise a client
  226. // already in JSONP mode gets a script-tag syntax error instead of the reason.
  227. func TestAuthErrorsHonorJSONP(t *testing.T) {
  228. m := newTestMiddleware(&stubValidator{})
  229. t.Run("session error", func(t *testing.T) {
  230. h := m.RequireSession(&stubSessionResolver{}, func(http.ResponseWriter, *http.Request, *Session) {
  231. t.Fatal("handler should not run")
  232. })
  233. r := httptest.NewRequest(http.MethodGet, "/im/sendIM?c=_callbacks_._x&r=7", nil)
  234. w := httptest.NewRecorder()
  235. h.ServeHTTP(w, r)
  236. body := w.Body.String()
  237. assert.True(t, strings.HasPrefix(body, "_callbacks_._x("), "got %s", body)
  238. assert.True(t, strings.HasSuffix(body, ");"), "got %s", body)
  239. assert.Contains(t, body, `"statusCode":400`)
  240. assert.Contains(t, body, `"requestId":"7"`)
  241. // A 4xx would stop the browser executing the script tag.
  242. assert.Equal(t, http.StatusOK, w.Code)
  243. })
  244. t.Run("missing credentials", func(t *testing.T) {
  245. h := m.AuthenticateFlexible(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
  246. t.Fatal("handler should not run")
  247. }))
  248. r := httptest.NewRequest(http.MethodGet, "/im/sendIM?c=cb&r=9", nil)
  249. w := httptest.NewRecorder()
  250. h.ServeHTTP(w, r)
  251. body := w.Body.String()
  252. assert.True(t, strings.HasPrefix(body, "cb("), "got %s", body)
  253. assert.Contains(t, body, `"statusCode":400`)
  254. assert.Contains(t, body, `"requestId":"9"`)
  255. assert.Equal(t, http.StatusOK, w.Code)
  256. })
  257. t.Run("without a callback the session error keeps its HTTP status", func(t *testing.T) {
  258. h := m.RequireSession(&stubSessionResolver{}, func(http.ResponseWriter, *http.Request, *Session) {
  259. t.Fatal("handler should not run")
  260. })
  261. r := httptest.NewRequest(http.MethodGet, "/im/sendIM", nil)
  262. w := httptest.NewRecorder()
  263. h.ServeHTTP(w, r)
  264. assert.Equal(t, http.StatusBadRequest, w.Code)
  265. assert.Contains(t, w.Header().Get("Content-Type"), "json")
  266. })
  267. }
  268. // An XML client cannot parse a JSON error, so it reports an unreadable response
  269. // instead of the reason the auth layer rejected it.
  270. func TestAuthErrorsHonorXML(t *testing.T) {
  271. m := newTestMiddleware(&stubValidator{})
  272. h := m.Authenticate(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
  273. t.Fatal("handler should not run")
  274. }))
  275. t.Run("format in the query string", func(t *testing.T) {
  276. r := httptest.NewRequest(http.MethodGet, "/aim/startOSCARSession?f=xml", nil)
  277. w := httptest.NewRecorder()
  278. h.ServeHTTP(w, r)
  279. assert.Contains(t, w.Header().Get("Content-Type"), "xml")
  280. assert.Contains(t, w.Body.String(), "<statusCode>400</statusCode>")
  281. })
  282. // A POST states the format in its body, the only place clientLogin sends it.
  283. t.Run("format in the POST body", func(t *testing.T) {
  284. r := httptest.NewRequest(http.MethodPost, "/auth/clientLogin", strings.NewReader("s=testuser&f=xml"))
  285. r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
  286. w := httptest.NewRecorder()
  287. h.ServeHTTP(w, r)
  288. assert.Contains(t, w.Header().Get("Content-Type"), "xml")
  289. assert.Contains(t, w.Body.String(), "<statusCode>400</statusCode>")
  290. })
  291. // A client on the <script> transport needs executable JS back whatever "f"
  292. // says; XML there is a script load failure with no reason attached.
  293. t.Run("a callback outranks the format", func(t *testing.T) {
  294. r := httptest.NewRequest(http.MethodGet, "/aim/startOSCARSession?f=xml&c=cb&r=3", nil)
  295. w := httptest.NewRecorder()
  296. h.ServeHTTP(w, r)
  297. body := w.Body.String()
  298. assert.True(t, strings.HasPrefix(body, "cb("), "got %s", body)
  299. assert.Contains(t, body, `"statusCode":400`)
  300. assert.Contains(t, body, `"requestId":"3"`)
  301. })
  302. }
  303. // stubSessionResolver never resolves a session, so RequireSession always rejects.
  304. type stubSessionResolver struct{}
  305. func (stubSessionResolver) GetSession(context.Context, string) (*Session, error) {
  306. return nil, assert.AnError
  307. }
  308. func (stubSessionResolver) TouchSession(context.Context, string) error { return nil }
  309. func newTestRateLimitMiddleware() *RateLimitMiddleware {
  310. return NewRateLimitMiddleware(wire.DefaultSNACRateLimits(), slog.New(slog.DiscardHandler))
  311. }
  312. // assertRateLimited checks that a response is the Web API's rate limit
  313. // rejection: HTTP 200 at the transport level so the client parses the body,
  314. // with envelope code 430 carrying the rejection inside. wantRetryAfter is the
  315. // expected Retry-After header, "" for a rejection that carries none.
  316. func assertRateLimited(t *testing.T, rec *httptest.ResponseRecorder, wantRetryAfter string) {
  317. t.Helper()
  318. assert.Equal(t, http.StatusOK, rec.Code)
  319. assert.Equal(t, wantRetryAfter, rec.Header().Get("Retry-After"))
  320. var envelope struct {
  321. Response struct {
  322. StatusCode int `json:"statusCode"`
  323. StatusText string `json:"statusText"`
  324. } `json:"response"`
  325. }
  326. assert.NoError(t, json.Unmarshal(rec.Body.Bytes(), &envelope))
  327. assert.Equal(t, statusRateLimited, envelope.Response.StatusCode)
  328. assert.Equal(t, "rate limit exceeded", envelope.Response.StatusText)
  329. }
  330. func TestRateLimitMiddleware_OSCAR(t *testing.T) {
  331. tests := []struct {
  332. name string
  333. // foodGroup/subGroup passed to the middleware
  334. foodGroup uint16
  335. subGroup uint16
  336. // requests is how many times the wrapped handler is invoked
  337. requests int
  338. // wantCalls is how many of those requests reach the handler
  339. wantCalls int
  340. // wantRetryAfter is the Retry-After header on the final rejection, "" for
  341. // a rejection that carries none
  342. wantRetryAfter string
  343. }{
  344. {
  345. name: "first request is allowed",
  346. foodGroup: wire.ICBM,
  347. subGroup: wire.ICBMChannelMsgToHost,
  348. requests: 1,
  349. wantCalls: 1,
  350. },
  351. {
  352. name: "a burst trips the limit",
  353. foodGroup: wire.ICBM,
  354. subGroup: wire.ICBMChannelMsgToHost,
  355. requests: 5,
  356. wantCalls: 1,
  357. // retryAfterFor rounds the sub-second wait these tight classes need
  358. // up to the minRetryAfter floor.
  359. wantRetryAfter: "1",
  360. },
  361. {
  362. name: "sustained abuse escalates to disconnect",
  363. foodGroup: wire.ICBM,
  364. subGroup: wire.ICBMChannelMsgToHost,
  365. // the 7th request drives the average below DisconnectLevel
  366. requests: 7,
  367. wantCalls: 1,
  368. // a disconnected session has no aimsid left to retry with
  369. wantRetryAfter: "",
  370. },
  371. {
  372. // A non-IM class is enforced the same way. The middleware only ever
  373. // rejects; notifying the client is the per-account monitor's job, and
  374. // it surfaces only the IM class to the alert.
  375. name: "a non-IM class is still enforced",
  376. foodGroup: wire.Feedbag,
  377. subGroup: wire.FeedbagInsertItem,
  378. requests: 5,
  379. wantCalls: 1,
  380. wantRetryAfter: "1",
  381. },
  382. {
  383. name: "unmapped SNAC fails open",
  384. // 0xFFFF is not a food group, so no rate class maps to it
  385. foodGroup: 0xFFFF,
  386. subGroup: 0xFFFF,
  387. requests: 5,
  388. wantCalls: 5,
  389. },
  390. }
  391. for _, tt := range tests {
  392. t.Run(tt.name, func(t *testing.T) {
  393. session := newTestWebAPISession(t, tightRateLimitClasses())
  394. middleware := newTestRateLimitMiddleware()
  395. calls := 0
  396. handler := middleware.OSCAR(tt.foodGroup, tt.subGroup)(
  397. func(w http.ResponseWriter, r *http.Request, s *Session) {
  398. calls++
  399. w.WriteHeader(http.StatusOK)
  400. })
  401. var last *httptest.ResponseRecorder
  402. for range tt.requests {
  403. last = httptest.NewRecorder()
  404. handler(last, httptest.NewRequest(http.MethodGet, "/im/sendIM", nil), session)
  405. }
  406. assert.Equal(t, tt.wantCalls, calls)
  407. // The middleware enforces the limit but never pushes rate limit
  408. // events; that is the per-account monitor's responsibility.
  409. assert.Empty(t, rateLimitEventStatuses(t, session))
  410. if tt.wantCalls < tt.requests {
  411. assertRateLimited(t, last, tt.wantRetryAfter)
  412. } else {
  413. assert.Equal(t, http.StatusOK, last.Code)
  414. }
  415. })
  416. }
  417. }
  418. // Retry-After must name a wait that actually clears the limit. A rejected
  419. // request is still charged, so a client retrying on the advertised interval
  420. // drives the moving average toward that interval: a hint below the class's
  421. // ClearLevel holds the average under the bar and the client stays limited
  422. // forever.
  423. func TestRetryAfterFor_clearsTheLimit(t *testing.T) {
  424. classes := wire.DefaultRateLimitClasses()
  425. for _, class := range classes.All() {
  426. t.Run(fmt.Sprintf("class %d", class.ID), func(t *testing.T) {
  427. sess := state.NewSession()
  428. sess.AddInstance()
  429. now := time.Now()
  430. sess.SetRateClasses(now, classes)
  431. // Burst until the class trips.
  432. var status wire.RateLimitStatus
  433. for status != wire.RateLimitStatusLimited {
  434. now = now.Add(100 * time.Millisecond)
  435. status = sess.EvaluateRateLimit(now, class.ID)
  436. require.NotEqual(t, wire.RateLimitStatusDisconnect, status, "burst escalated past limited")
  437. }
  438. // Wait exactly as long as the rejection advertised, then retry.
  439. retryAfter := retryAfterFor(sess.RateLimitStates()[class.ID-1])
  440. now = now.Add(retryAfter)
  441. assert.Equal(t, wire.RateLimitStatusClear, sess.EvaluateRateLimit(now, class.ID),
  442. "honoring a Retry-After of %s did not clear the limit", retryAfter)
  443. })
  444. }
  445. }
  446. // The rejection advertises the wait computed for the class it charged, not a
  447. // flat one.
  448. func TestRateLimitMiddleware_OSCAR_retryAfterMatchesClass(t *testing.T) {
  449. session := newTestWebAPISession(t, wire.DefaultRateLimitClasses())
  450. middleware := newTestRateLimitMiddleware()
  451. handler := middleware.OSCAR(wire.ICBM, wire.ICBMChannelMsgToHost)(
  452. func(w http.ResponseWriter, r *http.Request, s *Session) {})
  453. classID, ok := wire.DefaultSNACRateLimits().RateClassLookup(wire.ICBM, wire.ICBMChannelMsgToHost)
  454. require.True(t, ok)
  455. // Back-to-back requests trip the limit; keep going until one is rejected.
  456. var rec *httptest.ResponseRecorder
  457. for range 20 {
  458. rec = httptest.NewRecorder()
  459. handler(rec, httptest.NewRequest(http.MethodGet, "/im/sendIM", nil), session)
  460. if rec.Header().Get("Retry-After") != "" {
  461. break
  462. }
  463. }
  464. want := retryAfterFor(session.OSCARSession.Session().RateLimitStates()[classID-1])
  465. assert.Equal(t, fmt.Sprintf("%d", int(want.Seconds())), rec.Header().Get("Retry-After"))
  466. // The production IM class clears at 5100ms, so the wait is necessarily
  467. // longer than the flat 5s hint this replaced.
  468. assert.Greater(t, want, 5*time.Second)
  469. }
  470. // The rejection is encoded in whatever format the request negotiates, via the
  471. // same SendResponse path a normal handler uses.
  472. func TestRateLimitMiddleware_sendRateLimited(t *testing.T) {
  473. tests := []struct {
  474. name string
  475. // query is appended to the request URL
  476. query string
  477. // wantCode is the transport status
  478. wantCode int
  479. // wantBody is the exact response body
  480. wantBody string
  481. // wantContentType is a substring of the Content-Type header
  482. wantContentType string
  483. }{
  484. {
  485. name: "plain JSON",
  486. query: "",
  487. wantCode: http.StatusOK,
  488. wantBody: `{"response":{"statusCode":430,"statusText":"rate limit exceeded","data":{}}}`,
  489. wantContentType: "application/json",
  490. },
  491. {
  492. name: "JSONP callback",
  493. query: "?c=myCallback",
  494. wantCode: http.StatusOK,
  495. wantBody: `myCallback({"response":{"statusCode":430,"statusText":"rate limit exceeded","data":{}}});`,
  496. wantContentType: "application/javascript",
  497. },
  498. {
  499. // The client correlates a JSONP reply solely by response.requestId,
  500. // so the request's "r" param must be echoed back or the request hangs.
  501. name: "JSONP callback echoes requestId",
  502. query: "?c=myCallback&r=42",
  503. wantCode: http.StatusOK,
  504. wantBody: `myCallback({"response":{"statusCode":430,"statusText":"rate limit exceeded","requestId":"42","data":{}}});`,
  505. wantContentType: "application/javascript",
  506. },
  507. {
  508. // Parens would let the callback name inject script; SendResponse
  509. // rejects the malformed callback rather than reflecting it.
  510. name: "invalid JSONP callback is rejected",
  511. query: "?c=alert(1)",
  512. wantCode: http.StatusBadRequest,
  513. // sendJSONError encodes with json.Encoder, which appends a newline.
  514. wantBody: "{\"response\":{\"statusCode\":400,\"statusText\":\"invalid callback parameter\",\"data\":{}}}\n",
  515. wantContentType: "application/json",
  516. },
  517. }
  518. for _, tt := range tests {
  519. t.Run(tt.name, func(t *testing.T) {
  520. rec := httptest.NewRecorder()
  521. newTestRateLimitMiddleware().sendRateLimited(rec, httptest.NewRequest(http.MethodGet, "/im/sendIM"+tt.query, nil), 5*time.Second)
  522. body, err := io.ReadAll(rec.Body)
  523. assert.NoError(t, err)
  524. assert.Equal(t, tt.wantCode, rec.Code)
  525. assert.Contains(t, rec.Header().Get("Content-Type"), tt.wantContentType)
  526. assert.Equal(t, tt.wantBody, string(body))
  527. })
  528. }
  529. }
  530. // A client asking for XML or AMF (via f= or the Accept header) still gets a
  531. // parseable rejection envelope rather than JSON, since the rejection rides the
  532. // same SendResponse path as a normal handler.
  533. func TestRateLimitMiddleware_sendRateLimited_nonJSONFormats(t *testing.T) {
  534. t.Run("f=xml", func(t *testing.T) {
  535. rec := httptest.NewRecorder()
  536. newTestRateLimitMiddleware().sendRateLimited(rec, httptest.NewRequest(http.MethodGet, "/im/sendIM?f=xml", nil), 5*time.Second)
  537. assert.Equal(t, http.StatusOK, rec.Code)
  538. assert.Equal(t, "5", rec.Header().Get("Retry-After"))
  539. assert.Contains(t, rec.Header().Get("Content-Type"), "xml")
  540. body := rec.Body.String()
  541. assert.Contains(t, body, "<statusCode>430</statusCode>")
  542. assert.Contains(t, body, "rate limit exceeded")
  543. })
  544. t.Run("f=amf", func(t *testing.T) {
  545. rec := httptest.NewRecorder()
  546. newTestRateLimitMiddleware().sendRateLimited(rec, httptest.NewRequest(http.MethodGet, "/im/sendIM?f=amf", nil), 5*time.Second)
  547. assert.Equal(t, http.StatusOK, rec.Code)
  548. assert.Equal(t, "5", rec.Header().Get("Retry-After"))
  549. assert.Contains(t, rec.Header().Get("Content-Type"), "amf")
  550. assert.NotEmpty(t, rec.Body.Bytes())
  551. })
  552. t.Run("Accept amf", func(t *testing.T) {
  553. rec := httptest.NewRecorder()
  554. req := httptest.NewRequest(http.MethodGet, "/im/sendIM", nil)
  555. req.Header.Set("Accept", "application/x-amf")
  556. newTestRateLimitMiddleware().sendRateLimited(rec, req, 5*time.Second)
  557. assert.Equal(t, http.StatusOK, rec.Code)
  558. assert.Contains(t, rec.Header().Get("Content-Type"), "amf")
  559. assert.NotEmpty(t, rec.Body.Bytes())
  560. })
  561. }
  562. func TestRateLimitStatusName(t *testing.T) {
  563. tests := []struct {
  564. name string
  565. status wire.RateLimitStatus
  566. want string
  567. }{
  568. {name: "clear", status: wire.RateLimitStatusClear, want: "clear"},
  569. {name: "alert maps to the client's warn", status: wire.RateLimitStatusAlert, want: "warn"},
  570. {name: "limited", status: wire.RateLimitStatusLimited, want: "limit"},
  571. {name: "disconnect", status: wire.RateLimitStatusDisconnect, want: "disconnect"},
  572. {name: "unknown status has no client equivalent", status: wire.RateLimitStatus(0), want: ""},
  573. }
  574. for _, tt := range tests {
  575. t.Run(tt.name, func(t *testing.T) {
  576. assert.Equal(t, tt.want, rateLimitStatusName(tt.status))
  577. })
  578. }
  579. }