tests: add /debug endpoint to fakestore to allow interrupting downloads

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,178 @@ type detailsReplyJSON struct {
 	Base            string   `json:"base,omitempty"`
 }
 
+type killAfterWriter struct {
+	http.ResponseWriter
+	killAfter int64
+}
+
+func (kaw *killAfterWriter) Write(p []byte) (int, error) {
+	if kaw.killAfter <= 0 {
+		// already exceeded the quota, kill immediately
+		kaw.hijackAndClose()
+		return 0, fmt.Errorf("connection killed")
+	}
+
+	toWrite := p
+	shouldKill := false
+	if int64(len(p)) > kaw.killAfter {
+		// write only up to the remaining quota
+		toWrite = p[:kaw.killAfter]
+		shouldKill = true
+	}
+
+	n, err := kaw.ResponseWriter.Write(toWrite)
+	kaw.killAfter -= int64(n)
+
+	if shouldKill {
+		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
+
+		killAfter, exists := func() (int64, bool) {
+			s.lock.Lock()
+			defer s.lock.Unlock()
+			v, ok := s.killAfter[path]
+			return v, ok
+		}()
+
+		if !exists {
+			handler(w, req)
+			return
+		}
+
+		kaw := &killAfterWriter{
+			ResponseWriter: w,
+			killAfter:      killAfter,
+		}
+		handler(kaw, req)
+
+		if kaw.killAfter < 0 {
+			logger.Noticef("%s was force killed, quota exceeded", path)
+		}
+
+		// update killAfter for path after write finishes
+		func() {
+			s.lock.Lock()
+			defer s.lock.Unlock()
+			s.killAfter[path] = kaw.killAfter
+		}()
+	}
+}
+
 func (s *Store) searchEndpoint(w http.ResponseWriter, req *http.Request) {
 	w.WriteHeader(501)
 	fmt.Fprintf(w, "search not implemented")

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":{}}`)
+}