diff --git a/tests/lib/fakestore/store/store.go b/tests/lib/fakestore/store/store.go index 8d4cc707cdf..03d13111004 100644 --- a/tests/lib/fakestore/store/store.go +++ b/tests/lib/fakestore/store/store.go @@ -35,6 +35,7 @@ import ( "regexp" "strconv" "strings" + "sync" "time" "github.com/snapcore/snapd/asserts" @@ -67,6 +68,8 @@ type snapCachedInfo struct { // Store is our snappy software store implementation type Store struct { + lock sync.Mutex + url string blobDir string assertDir string @@ -79,6 +82,11 @@ type Store struct { channelRepository *ChannelRepository snapsCache map[string]snapCachedInfo + + // endpoint -> quota value, note this is stateful, i.e. the quota is counted + // for all requests to a given endpoint and after exceeding it, all + // subsequent requests will fail until it is reset through a request + killAfter map[string]int64 } // NewStore creates a new store server serving snaps from the given top directory and assertions from topDir/asserts. If assertFallback is true missing assertions are looked up in the main online store. @@ -105,13 +113,16 @@ func NewStore(topDir, addr string, assertFallback bool) *Store { rootDir: filepath.Join(topDir, "channels"), }, snapsCache: make(map[string]snapCachedInfo), + killAfter: make(map[string]int64), } mux.HandleFunc("/", rootEndpoint) mux.HandleFunc("/api/v1/snaps/search", store.searchEndpoint) mux.HandleFunc("/api/v1/snaps/details/", store.detailsEndpoint) mux.HandleFunc("/api/v1/snaps/metadata", store.bulkEndpoint) - mux.Handle("/download/", http.StripPrefix("/download/", http.FileServer(http.Dir(topDir)))) + + fileServer := http.StripPrefix("/download/", http.FileServer(http.Dir(topDir))) + mux.Handle("/download/", logRangeHeader(store.applyKillAfter(fileServer.ServeHTTP))) mux.HandleFunc("/api/v1/snaps/auth/nonces", store.nonceEndpoint) mux.HandleFunc("/api/v1/snaps/auth/sessions", store.sessionEndpoint) @@ -122,6 +133,8 @@ func NewStore(topDir, addr string, assertFallback bool) *Store { mux.HandleFunc("/v2/repairs/", store.repairsEndpoint) + mux.HandleFunc("/debug", store.debugEndpoint) + return store } @@ -354,6 +367,189 @@ type detailsReplyJSON struct { Base string `json:"base,omitempty"` } +type killAfterWriter struct { + http.ResponseWriter + path string + consumeQuota func(want int) int +} + +func (kaw *killAfterWriter) Write(p []byte) (int, error) { + toWrite := p + shouldKill := false + + got := kaw.consumeQuota(len(toWrite)) + if len(p) > got { + // write only up to the remaining quota + toWrite = p[:got] + shouldKill = true + } + + n, err := kaw.ResponseWriter.Write(toWrite) + + if shouldKill { + logger.Noticef("request to %s was force killed, quota exceeded", kaw.path) + kaw.hijackAndClose() + return n, fmt.Errorf("connection killed") + } + + return n, err +} + +func (kaw *killAfterWriter) hijackAndClose() { + // flush any buffered data before closing + if f, ok := kaw.ResponseWriter.(http.Flusher); ok { + f.Flush() + } + // and proceed to close + hj, ok := kaw.ResponseWriter.(http.Hijacker) + if ok { + conn, _, _ := hj.Hijack() + conn.Close() + } +} + +type debugRequestJSON struct { + Action string `json:"action"` + + KillPath string `json:"kill-path"` + KillAfter int64 `json:"kill-after"` +} + +type debugResultJSON struct { + KillAfter map[string]int64 `json:"kill-after"` +} + +func (s *Store) debugEndpoint(w http.ResponseWriter, req *http.Request) { + if req.Method == http.MethodGet { + out, err := func() ([]byte, error) { + s.lock.Lock() + defer s.lock.Unlock() + res := debugResultJSON{ + KillAfter: s.killAfter, + } + return json.Marshal(res) + }() + if err != nil { + http.Error(w, fmt.Sprintf("cannot marshal: %v", err), 500) + return + } + w.Write(out) + return + } + + if req.Method != http.MethodPost { + w.WriteHeader(405) // Method Not Allowed + return + } + + var debugReq *debugRequestJSON + decoder := json.NewDecoder(req.Body) + if err := decoder.Decode(&debugReq); err != nil { + http.Error(w, fmt.Sprintf("cannot decode request body: %v", err), 400) + return + } + + var err error + switch debugReq.Action { + case "kill-request": + err = s.debugActionKillDownload(debugReq) + case "reset": + s.debugActionReset(debugReq) + default: + err = fmt.Errorf("unexpected debug action %q", debugReq.Action) + } + if err != nil { + w.WriteHeader(400) + fmt.Fprint(w, err.Error()) + } +} + +func (s *Store) debugActionKillDownload(debugReq *debugRequestJSON) error { + if debugReq.KillPath == "" { + return fmt.Errorf("kill-path cannot be empty") + } + + if strings.HasPrefix(debugReq.KillPath, "/debug/") { + return fmt.Errorf("kill-path cannot be applied to /debug/ endpoints") + } + + s.lock.Lock() + defer s.lock.Unlock() + + if debugReq.KillAfter == 0 { + delete(s.killAfter, debugReq.KillPath) + } else { + s.killAfter[debugReq.KillPath] = debugReq.KillAfter + } + return nil +} + +func (s *Store) debugActionReset(debugReq *debugRequestJSON) { + s.lock.Lock() + defer s.lock.Unlock() + + s.killAfter = map[string]int64{} +} + +func logRangeHeader(handler http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, req *http.Request) { + path := req.URL.Path + if len(req.Header["Range"]) > 0 { + logger.Noticef(`requested range for %s is %v`, path, req.Header["Range"]) + } + handler(w, req) + } +} + +func (s *Store) applyKillAfter(handler http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, req *http.Request) { + path := req.URL.Path + + exists := func() bool { + s.lock.Lock() + defer s.lock.Unlock() + _, ok := s.killAfter[path] + return ok + }() + + if !exists { + handler(w, req) + return + } + + kaw := &killAfterWriter{ + ResponseWriter: w, + path: path, + consumeQuota: func(want int) int { + s.lock.Lock() + defer s.lock.Unlock() + + v, ok := s.killAfter[path] + if !ok { + // no quota set + return want + } + + left := int(v) + + var got int + if want > left { + got = left + left = 0 + } else { + got = want + left -= want + } + s.killAfter[path] = int64(left) + + return got + }, + } + handler(kaw, req) + + } +} + func (s *Store) searchEndpoint(w http.ResponseWriter, req *http.Request) { w.WriteHeader(501) fmt.Fprintf(w, "search not implemented") diff --git a/tests/lib/fakestore/store/store_test.go b/tests/lib/fakestore/store/store_test.go index 3534ba7bdda..05e082ca074 100644 --- a/tests/lib/fakestore/store/store_test.go +++ b/tests/lib/fakestore/store/store_test.go @@ -1303,3 +1303,142 @@ func (s *storeTestSuite) TestSnapActionEndpointUnknownSnapAutoRefresh(c *C) { }, }) } + +func (s *storeTestSuite) TestDebugEndpointKillAfter(c *C) { + snapFn := s.makeTestSnap(c, "name: foo\nversion: 1") + snapInfo, err := os.Stat(snapFn) + c.Assert(err, IsNil) + + downloadPath := "/download/foo_1_all.snap" + killAfter := int64(512) + c.Assert(snapInfo.Size() > killAfter, Equals, true, + Commentf("test snap must be larger than kill-after threshold")) + + // Set a rule + resp, err := s.StorePostJSON("/debug", []byte(fmt.Sprintf(`{ + "action": "kill-request", + "kill-path": "%s", + "kill-after": %d + }`, downloadPath, killAfter))) + c.Assert(err, IsNil) + resp.Body.Close() + + resp, err = s.StoreGet("/debug") + c.Assert(err, IsNil) + defer resp.Body.Close() + + c.Assert(resp.StatusCode, Equals, 200) + var body debugResultJSON + c.Assert(json.NewDecoder(resp.Body).Decode(&body), IsNil) + c.Check(body.KillAfter, DeepEquals, map[string]int64{ + downloadPath: killAfter, + }) + + // Download is interrupted, we get fewer bytes than the full snap + resp, err = s.StoreGet(downloadPath) + c.Assert(err, IsNil) + defer resp.Body.Close() + + got, _ := io.ReadAll(resp.Body) + // Connection forcefully closed mid-transfer, exactly killAfter bytes received + c.Check(int64(len(got)), Equals, killAfter) + + // Retry the request, which should be killed after receiving 0 bytes because + // the killAfter effect is stateful. + resp, err = s.StoreGet(downloadPath) + c.Assert(err, IsNil) + defer resp.Body.Close() + + got, _ = io.ReadAll(resp.Body) + // Connection forcefully closed mid-transfer, exactly killAfter bytes received + c.Check(int64(len(got)), Equals, int64(0)) + + // Clear it by setting kill-after to 0 + resp, err = s.StorePostJSON("/debug", []byte(fmt.Sprintf(`{ + "action": "kill-request", + "kill-path": "%s", + "kill-after": 0 + }`, downloadPath))) + c.Assert(err, IsNil) + resp.Body.Close() + + resp, err = s.StoreGet("/debug") + c.Assert(err, IsNil) + defer resp.Body.Close() + + var bodyAfterClear debugResultJSON + c.Assert(json.NewDecoder(resp.Body).Decode(&bodyAfterClear), IsNil) + c.Check(bodyAfterClear.KillAfter, HasLen, 0) + + // Download succeeds after clearing kill-after + resp, err = s.StoreGet(downloadPath) + c.Assert(err, IsNil) + defer resp.Body.Close() + + c.Assert(resp.StatusCode, Equals, 200) + got, err = io.ReadAll(resp.Body) + c.Assert(err, IsNil) + c.Check(int64(len(got)), Equals, snapInfo.Size()) +} + +func (s *storeTestSuite) TestDebugEndpointUnknownAction(c *C) { + resp, err := s.StorePostJSON("/debug", []byte(`{ + "action": "unknown-action" + }`)) + c.Assert(err, IsNil) + defer resp.Body.Close() + + c.Assert(resp.StatusCode, Equals, 400) + body, err := io.ReadAll(resp.Body) + c.Assert(err, IsNil) + c.Check(string(body), Equals, `unexpected debug action "unknown-action"`) +} + +func (s *storeTestSuite) TestDebugEndpointMethodNotAllowed(c *C) { + req, err := http.NewRequest(http.MethodPut, s.store.URL()+"/debug", nil) + c.Assert(err, IsNil) + resp, err := s.client.Do(req) + c.Assert(err, IsNil) + defer resp.Body.Close() + + c.Assert(resp.StatusCode, Equals, 405) +} + +func (s *storeTestSuite) TestDebugActionReset(c *C) { + // Set a rule for endpoint connection interrupt + resp, err := s.StorePostJSON("/debug", []byte(`{ + "action": "kill-request", + "kill-path": "/foo/bar", + "kill-after": 123 + }`)) + c.Assert(err, IsNil) + resp.Body.Close() + c.Assert(resp.StatusCode, Equals, 200) + + resp, err = s.StoreGet("/debug") + c.Assert(err, IsNil) + defer resp.Body.Close() + + var buf bytes.Buffer + c.Assert(resp.StatusCode, Equals, 200) + _, err = io.Copy(&buf, resp.Body) + c.Assert(err, IsNil) + c.Check(buf.String(), Equals, `{"kill-after":{"/foo/bar":123}}`) + + // Clear it by setting kill-after to 0 + resp, err = s.StorePostJSON("/debug", []byte(`{ + "action": "reset" + }`)) + c.Assert(err, IsNil) + resp.Body.Close() + c.Assert(resp.StatusCode, Equals, 200) + + resp, err = s.StoreGet("/debug") + c.Assert(err, IsNil) + defer resp.Body.Close() + + buf.Reset() + _, err = io.Copy(&buf, resp.Body) + c.Assert(err, IsNil) + c.Check(buf.String(), Equals, `{"kill-after":{}}`) +}