4
0

middleware_test.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689
  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. func newTestMiddleware() *AuthMiddleware {
  19. return NewAuthMiddleware(discardLogger())
  20. }
  21. // The Web AIM client permanently downgrades to JSONP when a cross-origin
  22. // response arrives without Access-Control-Allow-Origin: it reads the response
  23. // the browser blocked as a status-0 empty one, and aim.client.js onXhrFailed_
  24. // clears its useXhr flag and never sets it again. One uncovered response is
  25. // enough to latch that, so the CORS layer wraps the whole mux and every layer
  26. // below it — routed, rejected, and unrouted alike — answers with the header.
  27. //
  28. // These cases go through the handler the listener actually serves, so they
  29. // cover the configured allowlist end to end rather than the options struct.
  30. func TestServer_CORS(t *testing.T) {
  31. // What most cases are configured with; a case overrides it to exercise the
  32. // wildcard and empty modes.
  33. defaultOrigins := []string{"https://ras.dev", "http://ras.dev", "http://localhost:8000"}
  34. tests := []struct {
  35. name string
  36. // nil means defaultOrigins.
  37. allowedOrigins []string
  38. method string
  39. path string
  40. origin string
  41. // Setting reqMethod makes the request a preflight.
  42. reqMethod string
  43. reqHeaders string
  44. wantStatus int
  45. wantOrigin string
  46. wantMethods string
  47. wantHeaders string
  48. wantMaxAge string
  49. }{
  50. // Each layer that can produce a response must carry the header.
  51. {
  52. name: "public auth endpoint", method: http.MethodPost, path: "/auth/clientLogin",
  53. origin: "https://ras.dev", wantStatus: http.StatusBadRequest, wantOrigin: "https://ras.dev",
  54. },
  55. {
  56. name: "session rejection", method: http.MethodGet, path: "/aim/fetchEvents",
  57. origin: "https://ras.dev", wantStatus: http.StatusBadRequest, wantOrigin: "https://ras.dev",
  58. },
  59. {
  60. name: "stub route", method: http.MethodGet, path: "/aim/getData",
  61. origin: "https://ras.dev", wantStatus: http.StatusOK, wantOrigin: "https://ras.dev",
  62. },
  63. {
  64. name: "unrouted 404", method: http.MethodGet, path: "/service/getAttributes/nope",
  65. origin: "https://ras.dev", wantStatus: http.StatusNotFound, wantOrigin: "https://ras.dev",
  66. },
  67. // An origin is the scheme, host and port together, so the allowlist
  68. // admits exactly what it names and nothing adjacent.
  69. {
  70. name: "allowed origin, other scheme", method: http.MethodGet, path: "/aim/getData",
  71. origin: "http://ras.dev", wantStatus: http.StatusOK, wantOrigin: "http://ras.dev",
  72. },
  73. {
  74. name: "allowed origin, other port", method: http.MethodGet, path: "/aim/getData",
  75. origin: "http://localhost:8000", wantStatus: http.StatusOK, wantOrigin: "http://localhost:8000",
  76. },
  77. {
  78. name: "different port is a different origin", method: http.MethodGet, path: "/aim/getData",
  79. origin: "http://localhost", wantStatus: http.StatusOK,
  80. },
  81. {
  82. name: "different scheme is a different origin", method: http.MethodGet, path: "/aim/getData",
  83. origin: "https://localhost:8000", wantStatus: http.StatusOK,
  84. },
  85. {
  86. // The suffix a naive HasSuffix check would wave through.
  87. name: "attacker-registrable suffix", method: http.MethodGet, path: "/aim/getData",
  88. origin: "https://ras.dev.evil.example", wantStatus: http.StatusOK,
  89. },
  90. {
  91. name: "unlisted origin", method: http.MethodGet, path: "/aim/getData",
  92. origin: "http://evil.example", wantStatus: http.StatusOK,
  93. },
  94. {
  95. // A same-origin request sends no Origin, so there is nothing to allow.
  96. name: "no origin header", method: http.MethodGet, path: "/aim/getData",
  97. wantStatus: http.StatusOK,
  98. },
  99. // A preflight is answered by the CORS layer alone: it reaches no route,
  100. // so it needs no aimsid and 404s on no path.
  101. {
  102. name: "preflight GET", method: http.MethodOptions, path: "/im/sendIM",
  103. origin: "https://ras.dev", reqMethod: http.MethodGet,
  104. wantStatus: http.StatusNoContent, wantOrigin: "https://ras.dev",
  105. wantMethods: http.MethodGet, wantMaxAge: "3600",
  106. },
  107. {
  108. // Lowercase, as the Fetch spec requires a browser to send it: the
  109. // CORS layer matches these names byte for byte.
  110. name: "preflight POST with a header", method: http.MethodOptions, path: "/im/sendIM",
  111. origin: "https://ras.dev", reqMethod: http.MethodPost, reqHeaders: "content-type",
  112. wantStatus: http.StatusNoContent, wantOrigin: "https://ras.dev",
  113. wantMethods: http.MethodPost, wantHeaders: "content-type", wantMaxAge: "3600",
  114. },
  115. {
  116. name: "preflight on an unrouted path", method: http.MethodOptions, path: "/totally/unrouted",
  117. origin: "https://ras.dev", reqMethod: http.MethodPost,
  118. wantStatus: http.StatusNoContent, wantOrigin: "https://ras.dev",
  119. wantMethods: http.MethodPost, wantMaxAge: "3600",
  120. },
  121. {
  122. // Every route is a GET or a POST, so anything else is refused with
  123. // no CORS headers and the browser never sends the real request.
  124. name: "preflight PUT refused", method: http.MethodOptions, path: "/im/sendIM",
  125. origin: "https://ras.dev", reqMethod: http.MethodPut, wantStatus: http.StatusNoContent,
  126. },
  127. {
  128. name: "preflight DELETE refused", method: http.MethodOptions, path: "/im/sendIM",
  129. origin: "https://ras.dev", reqMethod: http.MethodDelete, wantStatus: http.StatusNoContent,
  130. },
  131. {
  132. name: "preflight with an unlisted header refused", method: http.MethodOptions, path: "/im/sendIM",
  133. origin: "https://ras.dev", reqMethod: http.MethodGet, reqHeaders: "x-not-allowed",
  134. wantStatus: http.StatusNoContent,
  135. },
  136. // The configured modes.
  137. {
  138. name: "wildcard allows any origin", allowedOrigins: []string{"*"},
  139. method: http.MethodGet, path: "/aim/getData", origin: "http://never.seen.example",
  140. wantStatus: http.StatusOK, wantOrigin: "http://never.seen.example",
  141. },
  142. {
  143. // envconfig splits on commas without trimming.
  144. name: "entries are trimmed", allowedOrigins: []string{" https://ras.dev ", " "},
  145. method: http.MethodGet, path: "/aim/getData", origin: "https://ras.dev",
  146. wantStatus: http.StatusOK, wantOrigin: "https://ras.dev",
  147. },
  148. {
  149. // An unset or empty value means "any origin", same as a lone *.
  150. name: "empty allowlist allows any origin", allowedOrigins: []string{" "},
  151. method: http.MethodGet, path: "/aim/getData", origin: "http://never.seen.example",
  152. wantStatus: http.StatusOK, wantOrigin: "http://never.seen.example",
  153. },
  154. {
  155. name: "unset allowlist allows any origin", allowedOrigins: []string{},
  156. method: http.MethodGet, path: "/aim/getData", origin: "http://never.seen.example",
  157. wantStatus: http.StatusOK, wantOrigin: "http://never.seen.example",
  158. },
  159. {
  160. // A wildcard may also stand in for the port.
  161. name: "port wildcard matches any port", allowedOrigins: []string{"http://localhost:*"},
  162. method: http.MethodGet, path: "/aim/getData", origin: "http://localhost:9999",
  163. wantStatus: http.StatusOK, wantOrigin: "http://localhost:9999",
  164. },
  165. {
  166. name: "port wildcard does not match another host", allowedOrigins: []string{"http://localhost:*"},
  167. method: http.MethodGet, path: "/aim/getData", origin: "http://evil.example:9999",
  168. wantStatus: http.StatusOK,
  169. },
  170. {
  171. // One wildcard per entry, standing in for 0 or more characters.
  172. name: "subdomain wildcard matches", allowedOrigins: []string{"https://*.example.com"},
  173. method: http.MethodGet, path: "/aim/getData", origin: "https://web.example.com",
  174. wantStatus: http.StatusOK, wantOrigin: "https://web.example.com",
  175. },
  176. {
  177. name: "subdomain wildcard does not match another domain", allowedOrigins: []string{"https://*.example.com"},
  178. method: http.MethodGet, path: "/aim/getData", origin: "https://web.evil.example",
  179. wantStatus: http.StatusOK,
  180. },
  181. }
  182. for _, tt := range tests {
  183. t.Run(tt.name, func(t *testing.T) {
  184. origins := tt.allowedOrigins
  185. if origins == nil {
  186. origins = defaultOrigins
  187. }
  188. h := testServerHandler(t, origins)
  189. r := httptest.NewRequest(tt.method, tt.path, nil)
  190. if tt.origin != "" {
  191. r.Header.Set("Origin", tt.origin)
  192. }
  193. if tt.reqMethod != "" {
  194. r.Header.Set("Access-Control-Request-Method", tt.reqMethod)
  195. }
  196. if tt.reqHeaders != "" {
  197. r.Header.Set("Access-Control-Request-Headers", tt.reqHeaders)
  198. }
  199. w := httptest.NewRecorder()
  200. h.ServeHTTP(w, r)
  201. assert.Equal(t, tt.wantStatus, w.Code)
  202. // An allowed origin is echoed back rather than answered with "*";
  203. // a rejected one gets no header at all, so the browser blocks it.
  204. assert.Equal(t, tt.wantOrigin, w.Header().Get("Access-Control-Allow-Origin"))
  205. assert.Contains(t, w.Header().Get("Vary"), "Origin")
  206. // Credentials are not enabled, so the browser is never told to send them.
  207. assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"))
  208. // These answer a preflight and mean nothing on any other response.
  209. assert.Equal(t, tt.wantMethods, w.Header().Get("Access-Control-Allow-Methods"))
  210. assert.Equal(t, tt.wantHeaders, strings.ToLower(w.Header().Get("Access-Control-Allow-Headers")))
  211. assert.Equal(t, tt.wantMaxAge, w.Header().Get("Access-Control-Max-Age"))
  212. })
  213. }
  214. }
  215. // testServerHandler builds the handler the listener actually serves, so the CORS
  216. // layer under test is the one NewServer configures.
  217. func testServerHandler(t *testing.T, allowedOrigins []string) http.Handler {
  218. t.Helper()
  219. handler := Handler{
  220. Logger: slog.Default(),
  221. AllowedOrigins: allowedOrigins,
  222. }
  223. srv := NewServer([]string{"127.0.0.1:0"}, slog.Default(), handler, NewSessionManager())
  224. require.NotEmpty(t, srv.servers)
  225. return srv.servers[0].Handler
  226. }
  227. func discardLogger() *slog.Logger {
  228. return slog.New(slog.NewTextHandler(io.Discard, nil))
  229. }
  230. func TestRequireSession_ReadsAimsidFromPOSTBody(t *testing.T) {
  231. tests := []struct {
  232. name string
  233. body string
  234. contentType string
  235. }{
  236. {
  237. name: "declared form body",
  238. body: "aimsid=abc&message=hello",
  239. contentType: "application/x-www-form-urlencoded",
  240. },
  241. {
  242. // No Content-Type announced; the request is form data all the same.
  243. name: "untyped form body",
  244. body: "aimsid=abc&message=hello",
  245. },
  246. }
  247. for _, tt := range tests {
  248. t.Run(tt.name, func(t *testing.T) {
  249. m := newTestMiddleware()
  250. reached := false
  251. h := m.RequireSession(okSessionResolver{},
  252. func(w http.ResponseWriter, r *http.Request, _ *Session) {
  253. reached = true
  254. assert.Equal(t, "hello", param(r, "message"))
  255. w.WriteHeader(http.StatusOK)
  256. })
  257. r := httptest.NewRequest(http.MethodPost, "/im/sendIM", strings.NewReader(tt.body))
  258. if tt.contentType != "" {
  259. r.Header.Set("Content-Type", tt.contentType)
  260. }
  261. w := httptest.NewRecorder()
  262. h.ServeHTTP(w, r)
  263. assert.True(t, reached, "request was rejected before reaching the handler")
  264. assert.Equal(t, http.StatusOK, w.Code)
  265. })
  266. }
  267. }
  268. // The session layer's own rejections must be JSONP-wrapped too, otherwise a
  269. // client already in JSONP mode gets a script-tag syntax error instead of the
  270. // reason.
  271. func TestAuthErrorsHonorJSONP(t *testing.T) {
  272. m := newTestMiddleware()
  273. t.Run("session error", func(t *testing.T) {
  274. h := m.RequireSession(&stubSessionResolver{}, func(http.ResponseWriter, *http.Request, *Session) {
  275. t.Fatal("handler should not run")
  276. })
  277. r := httptest.NewRequest(http.MethodGet, "/im/sendIM?c=_callbacks_._x&r=7", nil)
  278. w := httptest.NewRecorder()
  279. h.ServeHTTP(w, r)
  280. body := w.Body.String()
  281. assert.True(t, strings.HasPrefix(body, "_callbacks_._x("), "got %s", body)
  282. assert.True(t, strings.HasSuffix(body, ");"), "got %s", body)
  283. assert.Contains(t, body, `"statusCode":400`)
  284. assert.Contains(t, body, `"requestId":"7"`)
  285. // A 4xx would stop the browser executing the script tag.
  286. assert.Equal(t, http.StatusOK, w.Code)
  287. })
  288. t.Run("without a callback the session error keeps its HTTP status", func(t *testing.T) {
  289. h := m.RequireSession(&stubSessionResolver{}, func(http.ResponseWriter, *http.Request, *Session) {
  290. t.Fatal("handler should not run")
  291. })
  292. r := httptest.NewRequest(http.MethodGet, "/im/sendIM", nil)
  293. w := httptest.NewRecorder()
  294. h.ServeHTTP(w, r)
  295. assert.Equal(t, http.StatusBadRequest, w.Code)
  296. assert.Contains(t, w.Header().Get("Content-Type"), "json")
  297. })
  298. }
  299. // An XML client cannot parse a JSON error, so it reports an unreadable response
  300. // instead of the reason the session layer rejected it.
  301. func TestAuthErrorsHonorXML(t *testing.T) {
  302. m := newTestMiddleware()
  303. h := m.RequireSession(&stubSessionResolver{}, func(http.ResponseWriter, *http.Request, *Session) {
  304. t.Fatal("handler should not run")
  305. })
  306. t.Run("format in the query string", func(t *testing.T) {
  307. r := httptest.NewRequest(http.MethodGet, "/aim/startOSCARSession?f=xml", nil)
  308. w := httptest.NewRecorder()
  309. h.ServeHTTP(w, r)
  310. assert.Contains(t, w.Header().Get("Content-Type"), "xml")
  311. assert.Contains(t, w.Body.String(), "<statusCode>400</statusCode>")
  312. })
  313. // A POST states the format in its body, the only place clientLogin sends it.
  314. t.Run("format in the POST body", func(t *testing.T) {
  315. r := httptest.NewRequest(http.MethodPost, "/auth/clientLogin", strings.NewReader("s=testuser&f=xml"))
  316. r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
  317. w := httptest.NewRecorder()
  318. h.ServeHTTP(w, r)
  319. assert.Contains(t, w.Header().Get("Content-Type"), "xml")
  320. assert.Contains(t, w.Body.String(), "<statusCode>400</statusCode>")
  321. })
  322. // A client on the <script> transport needs executable JS back whatever "f"
  323. // says; XML there is a script load failure with no reason attached.
  324. t.Run("a callback outranks the format", func(t *testing.T) {
  325. r := httptest.NewRequest(http.MethodGet, "/aim/startOSCARSession?f=xml&c=cb&r=3", nil)
  326. w := httptest.NewRecorder()
  327. h.ServeHTTP(w, r)
  328. body := w.Body.String()
  329. assert.True(t, strings.HasPrefix(body, "cb("), "got %s", body)
  330. assert.Contains(t, body, `"statusCode":400`)
  331. assert.Contains(t, body, `"requestId":"3"`)
  332. })
  333. }
  334. // okSessionResolver always resolves, so RequireSession reaches the handler.
  335. type okSessionResolver struct{}
  336. func (okSessionResolver) GetSession(context.Context, string) (*Session, error) {
  337. return &Session{}, nil
  338. }
  339. func (okSessionResolver) TouchSession(context.Context, string) error { return nil }
  340. // stubSessionResolver never resolves a session, so RequireSession always rejects.
  341. type stubSessionResolver struct{}
  342. func (stubSessionResolver) GetSession(context.Context, string) (*Session, error) {
  343. return nil, assert.AnError
  344. }
  345. func (stubSessionResolver) TouchSession(context.Context, string) error { return nil }
  346. func newTestRateLimitMiddleware() *RateLimitMiddleware {
  347. return NewRateLimitMiddleware(wire.DefaultSNACRateLimits(), slog.New(slog.DiscardHandler))
  348. }
  349. // assertRateLimited checks that a response is the Web API's rate limit
  350. // rejection: HTTP 200 at the transport level so the client parses the body,
  351. // with envelope code 430 carrying the rejection inside. wantRetryAfter is the
  352. // expected Retry-After header, "" for a rejection that carries none.
  353. func assertRateLimited(t *testing.T, rec *httptest.ResponseRecorder, wantRetryAfter string) {
  354. t.Helper()
  355. assert.Equal(t, http.StatusOK, rec.Code)
  356. assert.Equal(t, wantRetryAfter, rec.Header().Get("Retry-After"))
  357. var envelope struct {
  358. Response struct {
  359. StatusCode int `json:"statusCode"`
  360. StatusText string `json:"statusText"`
  361. } `json:"response"`
  362. }
  363. assert.NoError(t, json.Unmarshal(rec.Body.Bytes(), &envelope))
  364. assert.Equal(t, statusRateLimited, envelope.Response.StatusCode)
  365. assert.Equal(t, "rate limit exceeded", envelope.Response.StatusText)
  366. }
  367. func TestRateLimitMiddleware_OSCAR(t *testing.T) {
  368. tests := []struct {
  369. name string
  370. // foodGroup/subGroup passed to the middleware
  371. foodGroup uint16
  372. subGroup uint16
  373. // requests is how many times the wrapped handler is invoked
  374. requests int
  375. // wantCalls is how many of those requests reach the handler
  376. wantCalls int
  377. // wantRetryAfter is the Retry-After header on the final rejection, "" for
  378. // a rejection that carries none
  379. wantRetryAfter string
  380. }{
  381. {
  382. name: "first request is allowed",
  383. foodGroup: wire.ICBM,
  384. subGroup: wire.ICBMChannelMsgToHost,
  385. requests: 1,
  386. wantCalls: 1,
  387. },
  388. {
  389. name: "a burst trips the limit",
  390. foodGroup: wire.ICBM,
  391. subGroup: wire.ICBMChannelMsgToHost,
  392. requests: 5,
  393. wantCalls: 1,
  394. // retryAfterFor rounds the sub-second wait these tight classes need
  395. // up to the minRetryAfter floor.
  396. wantRetryAfter: "1",
  397. },
  398. {
  399. name: "sustained abuse escalates to disconnect",
  400. foodGroup: wire.ICBM,
  401. subGroup: wire.ICBMChannelMsgToHost,
  402. // the 7th request drives the average below DisconnectLevel
  403. requests: 7,
  404. wantCalls: 1,
  405. // a disconnected session has no aimsid left to retry with
  406. wantRetryAfter: "",
  407. },
  408. {
  409. // A non-IM class is enforced the same way. The middleware only ever
  410. // rejects; notifying the client is the per-account monitor's job, and
  411. // it surfaces only the IM class to the alert.
  412. name: "a non-IM class is still enforced",
  413. foodGroup: wire.Feedbag,
  414. subGroup: wire.FeedbagInsertItem,
  415. requests: 5,
  416. wantCalls: 1,
  417. wantRetryAfter: "1",
  418. },
  419. {
  420. name: "unmapped SNAC fails open",
  421. // 0xFFFF is not a food group, so no rate class maps to it
  422. foodGroup: 0xFFFF,
  423. subGroup: 0xFFFF,
  424. requests: 5,
  425. wantCalls: 5,
  426. },
  427. }
  428. for _, tt := range tests {
  429. t.Run(tt.name, func(t *testing.T) {
  430. session := newTestWebAPISession(t, tightRateLimitClasses())
  431. middleware := newTestRateLimitMiddleware()
  432. calls := 0
  433. handler := middleware.OSCAR(tt.foodGroup, tt.subGroup)(
  434. func(w http.ResponseWriter, r *http.Request, s *Session) {
  435. calls++
  436. w.WriteHeader(http.StatusOK)
  437. })
  438. var last *httptest.ResponseRecorder
  439. for range tt.requests {
  440. last = httptest.NewRecorder()
  441. handler(last, httptest.NewRequest(http.MethodGet, "/im/sendIM", nil), session)
  442. }
  443. assert.Equal(t, tt.wantCalls, calls)
  444. // The middleware enforces the limit but never pushes rate limit
  445. // events; that is the per-account monitor's responsibility.
  446. assert.Empty(t, rateLimitEventStatuses(t, session))
  447. if tt.wantCalls < tt.requests {
  448. assertRateLimited(t, last, tt.wantRetryAfter)
  449. } else {
  450. assert.Equal(t, http.StatusOK, last.Code)
  451. }
  452. })
  453. }
  454. }
  455. // Retry-After must name a wait that actually clears the limit. A rejected
  456. // request is still charged, so a client retrying on the advertised interval
  457. // drives the moving average toward that interval: a hint below the class's
  458. // ClearLevel holds the average under the bar and the client stays limited
  459. // forever.
  460. func TestRetryAfterFor_clearsTheLimit(t *testing.T) {
  461. classes := wire.DefaultRateLimitClasses()
  462. for _, class := range classes.All() {
  463. t.Run(fmt.Sprintf("class %d", class.ID), func(t *testing.T) {
  464. sess := state.NewSession()
  465. sess.AddInstance()
  466. now := time.Now()
  467. sess.SetRateClasses(now, classes)
  468. // Burst until the class trips.
  469. var status wire.RateLimitStatus
  470. for status != wire.RateLimitStatusLimited {
  471. now = now.Add(100 * time.Millisecond)
  472. status = sess.EvaluateRateLimit(now, class.ID)
  473. require.NotEqual(t, wire.RateLimitStatusDisconnect, status, "burst escalated past limited")
  474. }
  475. // Wait exactly as long as the rejection advertised, then retry.
  476. retryAfter := retryAfterFor(sess.RateLimitStates()[class.ID-1])
  477. now = now.Add(retryAfter)
  478. assert.Equal(t, wire.RateLimitStatusClear, sess.EvaluateRateLimit(now, class.ID),
  479. "honoring a Retry-After of %s did not clear the limit", retryAfter)
  480. })
  481. }
  482. }
  483. // The rejection advertises the wait computed for the class it charged, not a
  484. // flat one.
  485. func TestRateLimitMiddleware_OSCAR_retryAfterMatchesClass(t *testing.T) {
  486. session := newTestWebAPISession(t, wire.DefaultRateLimitClasses())
  487. middleware := newTestRateLimitMiddleware()
  488. handler := middleware.OSCAR(wire.ICBM, wire.ICBMChannelMsgToHost)(
  489. func(w http.ResponseWriter, r *http.Request, s *Session) {})
  490. classID, ok := wire.DefaultSNACRateLimits().RateClassLookup(wire.ICBM, wire.ICBMChannelMsgToHost)
  491. require.True(t, ok)
  492. // Back-to-back requests trip the limit; keep going until one is rejected.
  493. var rec *httptest.ResponseRecorder
  494. for range 20 {
  495. rec = httptest.NewRecorder()
  496. handler(rec, httptest.NewRequest(http.MethodGet, "/im/sendIM", nil), session)
  497. if rec.Header().Get("Retry-After") != "" {
  498. break
  499. }
  500. }
  501. want := retryAfterFor(session.OSCARSession.Session().RateLimitStates()[classID-1])
  502. assert.Equal(t, fmt.Sprintf("%d", int(want.Seconds())), rec.Header().Get("Retry-After"))
  503. // The production IM class clears at 5100ms, so the wait is necessarily
  504. // longer than the flat 5s hint this replaced.
  505. assert.Greater(t, want, 5*time.Second)
  506. }
  507. // The rejection is encoded in whatever format the request negotiates, via the
  508. // same SendResponse path a normal handler uses.
  509. func TestRateLimitMiddleware_sendRateLimited(t *testing.T) {
  510. tests := []struct {
  511. name string
  512. // query is appended to the request URL
  513. query string
  514. // wantCode is the transport status
  515. wantCode int
  516. // wantBody is the exact response body
  517. wantBody string
  518. // wantContentType is a substring of the Content-Type header
  519. wantContentType string
  520. }{
  521. {
  522. name: "plain JSON",
  523. query: "",
  524. wantCode: http.StatusOK,
  525. wantBody: `{"response":{"statusCode":430,"statusText":"rate limit exceeded","data":{}}}`,
  526. wantContentType: "application/json",
  527. },
  528. {
  529. name: "JSONP callback",
  530. query: "?c=myCallback",
  531. wantCode: http.StatusOK,
  532. wantBody: `myCallback({"response":{"statusCode":430,"statusText":"rate limit exceeded","data":{}}});`,
  533. wantContentType: "application/javascript",
  534. },
  535. {
  536. // The client correlates a JSONP reply solely by response.requestId,
  537. // so the request's "r" param must be echoed back or the request hangs.
  538. name: "JSONP callback echoes requestId",
  539. query: "?c=myCallback&r=42",
  540. wantCode: http.StatusOK,
  541. wantBody: `myCallback({"response":{"statusCode":430,"statusText":"rate limit exceeded","requestId":"42","data":{}}});`,
  542. wantContentType: "application/javascript",
  543. },
  544. {
  545. // Parens would let the callback name inject script; SendResponse
  546. // rejects the malformed callback rather than reflecting it.
  547. name: "invalid JSONP callback is rejected",
  548. query: "?c=alert(1)",
  549. wantCode: http.StatusBadRequest,
  550. // sendJSONError encodes with json.Encoder, which appends a newline.
  551. wantBody: "{\"response\":{\"statusCode\":400,\"statusText\":\"invalid callback parameter\",\"data\":{}}}\n",
  552. wantContentType: "application/json",
  553. },
  554. }
  555. for _, tt := range tests {
  556. t.Run(tt.name, func(t *testing.T) {
  557. rec := httptest.NewRecorder()
  558. newTestRateLimitMiddleware().sendRateLimited(rec, httptest.NewRequest(http.MethodGet, "/im/sendIM"+tt.query, nil), 5*time.Second)
  559. body, err := io.ReadAll(rec.Body)
  560. assert.NoError(t, err)
  561. assert.Equal(t, tt.wantCode, rec.Code)
  562. assert.Contains(t, rec.Header().Get("Content-Type"), tt.wantContentType)
  563. assert.Equal(t, tt.wantBody, string(body))
  564. })
  565. }
  566. }
  567. // A client asking for XML or AMF (via f= or the Accept header) still gets a
  568. // parseable rejection envelope rather than JSON, since the rejection rides the
  569. // same SendResponse path as a normal handler.
  570. func TestRateLimitMiddleware_sendRateLimited_nonJSONFormats(t *testing.T) {
  571. t.Run("f=xml", func(t *testing.T) {
  572. rec := httptest.NewRecorder()
  573. newTestRateLimitMiddleware().sendRateLimited(rec, httptest.NewRequest(http.MethodGet, "/im/sendIM?f=xml", nil), 5*time.Second)
  574. assert.Equal(t, http.StatusOK, rec.Code)
  575. assert.Equal(t, "5", rec.Header().Get("Retry-After"))
  576. assert.Contains(t, rec.Header().Get("Content-Type"), "xml")
  577. body := rec.Body.String()
  578. assert.Contains(t, body, "<statusCode>430</statusCode>")
  579. assert.Contains(t, body, "rate limit exceeded")
  580. })
  581. t.Run("f=amf", func(t *testing.T) {
  582. rec := httptest.NewRecorder()
  583. newTestRateLimitMiddleware().sendRateLimited(rec, httptest.NewRequest(http.MethodGet, "/im/sendIM?f=amf", nil), 5*time.Second)
  584. assert.Equal(t, http.StatusOK, rec.Code)
  585. assert.Equal(t, "5", rec.Header().Get("Retry-After"))
  586. assert.Contains(t, rec.Header().Get("Content-Type"), "amf")
  587. assert.NotEmpty(t, rec.Body.Bytes())
  588. })
  589. t.Run("Accept amf", func(t *testing.T) {
  590. rec := httptest.NewRecorder()
  591. req := httptest.NewRequest(http.MethodGet, "/im/sendIM", nil)
  592. req.Header.Set("Accept", "application/x-amf")
  593. newTestRateLimitMiddleware().sendRateLimited(rec, req, 5*time.Second)
  594. assert.Equal(t, http.StatusOK, rec.Code)
  595. assert.Contains(t, rec.Header().Get("Content-Type"), "amf")
  596. assert.NotEmpty(t, rec.Body.Bytes())
  597. })
  598. }
  599. func TestRateLimitStatusName(t *testing.T) {
  600. tests := []struct {
  601. name string
  602. status wire.RateLimitStatus
  603. want string
  604. }{
  605. {name: "clear", status: wire.RateLimitStatusClear, want: "clear"},
  606. {name: "alert maps to the client's warn", status: wire.RateLimitStatusAlert, want: "warn"},
  607. {name: "limited", status: wire.RateLimitStatusLimited, want: "limit"},
  608. {name: "disconnect", status: wire.RateLimitStatusDisconnect, want: "disconnect"},
  609. {name: "unknown status has no client equivalent", status: wire.RateLimitStatus(0), want: ""},
  610. }
  611. for _, tt := range tests {
  612. t.Run(tt.name, func(t *testing.T) {
  613. assert.Equal(t, tt.want, rateLimitStatusName(tt.status))
  614. })
  615. }
  616. }