From b8faedc92fed1da2ce5fb0c26c9fe7f7301cd675 Mon Sep 17 00:00:00 2001 From: Douglas Pessoa Date: Mon, 17 Jun 2019 00:43:22 +0000 Subject: [PATCH] Add X-Remaining-Downloads and X-Remaining-Days headers to HEAD and GET responses --- server/handlers.go | 52 +++++++++++++++++++++++++++++++++------------- 1 file changed, 37 insertions(+), 15 deletions(-) diff --git a/server/handlers.go b/server/handlers.go index d88daed..4f16415 100644 --- a/server/handlers.go +++ b/server/handlers.go @@ -564,6 +564,14 @@ func getURL(r *http.Request) *url.URL { return u } +func calcRemainingLimits(metadata Metadata) (int, int) { + remainingDownloads := metadata.MaxDownloads - metadata.Downloads + timeDifference := metadata.MaxDate.Sub(time.Now()) + remainingDays := int(timeDifference.Hours()/24) + 1 + + return remainingDownloads, remainingDays +} + func (s *Server) Lock(token, filename string) error { key := path.Join(token, filename) @@ -583,7 +591,7 @@ func (s *Server) Unlock(token, filename string) error { return nil } -func (s *Server) CheckMetadata(token, filename string) error { +func (s *Server) CheckMetadata(token, filename string, increaseDownload bool) (Metadata, error) { s.Lock(token, filename) defer s.Unlock(token, filename) @@ -591,34 +599,36 @@ func (s *Server) CheckMetadata(token, filename string) error { r, _, _, err := s.storage.Get(token, fmt.Sprintf("%s.metadata", filename)) if s.storage.IsNotExist(err) { - return nil + return metadata, nil } else if err != nil { - return err + return metadata, err } defer r.Close() if err := json.NewDecoder(r).Decode(&metadata); err != nil { - return err + return metadata, err } else if metadata.Downloads >= metadata.MaxDownloads { - return errors.New("MaxDownloads expired.") + return metadata, errors.New("MaxDownloads expired.") } else if time.Now().After(metadata.MaxDate) { - return errors.New("MaxDate expired.") + return metadata, errors.New("MaxDate expired.") } else { // todo(nl5887): mutex? // update number of downloads - metadata.Downloads++ + if increaseDownload { + metadata.Downloads++ + } buffer := &bytes.Buffer{} if err := json.NewEncoder(buffer).Encode(metadata); err != nil { - return errors.New("Could not encode metadata") + return metadata, errors.New("Could not encode metadata") } else if err := s.storage.Put(token, fmt.Sprintf("%s.metadata", filename), buffer, "text/json", uint64(buffer.Len())); err != nil { - return errors.New("Could not save metadata") + return metadata, errors.New("Could not save metadata") } } - return nil + return metadata, nil } func (s *Server) CheckDeletionToken(deletionToken, token, filename string) error { @@ -688,7 +698,7 @@ func (s *Server) zipHandler(w http.ResponseWriter, r *http.Request) { token := strings.Split(key, "/")[0] filename := sanitize(strings.Split(key, "/")[1]) - if err := s.CheckMetadata(token, filename); err != nil { + if _, err := s.CheckMetadata(token, filename, true); err != nil { log.Printf("Error metadata: %s", err.Error()) continue } @@ -760,7 +770,7 @@ func (s *Server) tarGzHandler(w http.ResponseWriter, r *http.Request) { token := strings.Split(key, "/")[0] filename := sanitize(strings.Split(key, "/")[1]) - if err := s.CheckMetadata(token, filename); err != nil { + if _, err := s.CheckMetadata(token, filename, true); err != nil { log.Printf("Error metadata: %s", err.Error()) continue } @@ -819,7 +829,7 @@ func (s *Server) tarHandler(w http.ResponseWriter, r *http.Request) { token := strings.Split(key, "/")[0] filename := strings.Split(key, "/")[1] - if err := s.CheckMetadata(token, filename); err != nil { + if _, err := s.CheckMetadata(token, filename, true); err != nil { log.Printf("Error metadata: %s", err.Error()) continue } @@ -864,7 +874,9 @@ func (s *Server) headHandler(w http.ResponseWriter, r *http.Request) { token := vars["token"] filename := vars["filename"] - if err := s.CheckMetadata(token, filename); err != nil { + metadata, err := s.CheckMetadata(token, filename, false) + + if err != nil { log.Printf("Error metadata: %s", err.Error()) http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound) return @@ -880,9 +892,13 @@ func (s *Server) headHandler(w http.ResponseWriter, r *http.Request) { return } + remainingDownloads, remainingDays := calcRemainingLimits(metadata) + w.Header().Set("Content-Type", contentType) w.Header().Set("Content-Length", strconv.FormatUint(contentLength, 10)) w.Header().Set("Connection", "close") + w.Header().Set("X-Remaining-Downloads", strconv.Itoa(remainingDownloads)) + w.Header().Set("X-Remaining-Days", strconv.Itoa(remainingDays)) } func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) { @@ -892,7 +908,9 @@ func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) { token := vars["token"] filename := vars["filename"] - if err := s.CheckMetadata(token, filename); err != nil { + metadata, err := s.CheckMetadata(token, filename, true) + + if err != nil { log.Printf("Error metadata: %s", err.Error()) http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound) return @@ -918,10 +936,14 @@ func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) { disposition = "attachment" } + remainingDownloads, remainingDays := calcRemainingLimits(metadata) + w.Header().Set("Content-Type", contentType) w.Header().Set("Content-Length", strconv.FormatUint(contentLength, 10)) w.Header().Set("Content-Disposition", fmt.Sprintf("%s; filename=\"%s\"", disposition, filename)) w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Remaining-Downloads", strconv.Itoa(remainingDownloads)) + w.Header().Set("X-Remaining-Days", strconv.Itoa(remainingDays)) if w.Header().Get("Range") == "" { if _, err = io.Copy(w, reader); err != nil {