From 55e549e92c2fb70653741a18c7bc30cd15363681 Mon Sep 17 00:00:00 2001 From: Konrad Wojas Date: Sun, 3 May 2020 20:37:01 +0800 Subject: [PATCH] Move maxsizewriter to quota package --- handlers.go | 19 ------- maxsizewriter.go => quota/quota.go | 89 +++++++++++++++++++++--------- 2 files changed, 64 insertions(+), 44 deletions(-) rename maxsizewriter.go => quota/quota.go (52%) diff --git a/handlers.go b/handlers.go index 315fa4f..5b51e15 100644 --- a/handlers.go +++ b/handlers.go @@ -34,9 +34,6 @@ type Server struct { PrivateRepos bool Prometheus bool Debug bool - MaxRepoSize int64 - - repoSize int64 // must be accessed using sync/atomic } func (s *Server) isHashed(dir string) bool { @@ -487,22 +484,6 @@ func (s *Server) GetBlob(w http.ResponseWriter, r *http.Request) { } } -// tallySize counts the size of the contents of path. -func tallySize(path string) (int64, error) { - if path == "" { - path = "." - } - var size int64 - err := filepath.Walk(path, func(path string, info os.FileInfo, err error) error { - if err != nil { - return err - } - size += info.Size() - return nil - }) - return size, err -} - // SaveBlob saves a blob to the repository. func (s *Server) SaveBlob(w http.ResponseWriter, r *http.Request) { if s.Debug { diff --git a/maxsizewriter.go b/quota/quota.go similarity index 52% rename from maxsizewriter.go rename to quota/quota.go index 0258103..e037934 100644 --- a/maxsizewriter.go +++ b/quota/quota.go @@ -1,44 +1,67 @@ -package restserver +package quota import ( "fmt" "io" "net/http" + "os" + "path/filepath" "strconv" "sync/atomic" ) +// New creates a new quota Manager for given path. +// It will tally the current disk usage before returning. +func New(path string, maxSize int64) (*Manager, error) { + m := &Manager{ + path: path, + maxRepoSize: maxSize, + } + if err := m.updateSize(); err != nil { + return nil, err + } + return m, nil +} + +// Manager manages the repo quota for given filesystem root path, including subrepos +type Manager struct { + path string + maxRepoSize int64 + repoSize int64 // must be accessed using sync/atomic +} + // maxSizeWriter limits the number of bytes written // to the space that is currently available as given by // the server's MaxRepoSize. This type is safe for use // by multiple goroutines sharing the same *Server. type maxSizeWriter struct { io.Writer - server *Server + m *Manager } func (w maxSizeWriter) Write(p []byte) (n int, err error) { - if int64(len(p)) > w.server.repoSpaceRemaining() { - return 0, fmt.Errorf("repository has reached maximum size (%d bytes)", w.server.MaxRepoSize) + if int64(len(p)) > w.m.repoSpaceRemaining() { + return 0, fmt.Errorf("repository has reached maximum size (%d bytes)", w.m.maxRepoSize) } n, err = w.Writer.Write(p) - w.server.incrementRepoSpaceUsage(int64(n)) + w.m.incrementRepoSpaceUsage(int64(n)) return n, err } +func (m *Manager) updateSize() error { + // if we haven't yet computed the size of the repo, do so now + initialSize, err := tallySize(m.path) + if err != nil { + return err + } + atomic.StoreInt64(&m.repoSize, initialSize) + return nil +} + // maxSizeWriter wraps w in a writer that enforces s.MaxRepoSize. // If there is an error, a status code and the error are returned. -func (s *Server) maxSizeWriter(req *http.Request, w io.Writer) (io.Writer, int, error) { - // if we haven't yet computed the size of the repo, do so now - currentSize := atomic.LoadInt64(&s.repoSize) - if currentSize == 0 { - initialSize, err := tallySize(s.Path) - if err != nil { - return nil, http.StatusInternalServerError, err - } - atomic.StoreInt64(&s.repoSize, initialSize) - currentSize = initialSize - } +func (m *Manager) maxSizeWriter(req *http.Request, w io.Writer) (io.Writer, int, error) { + currentSize := atomic.LoadInt64(&m.repoSize) // if content-length is set and is trustworthy, we can save some time // and issue a polite error if it declares a size that's too big; since @@ -49,32 +72,48 @@ func (s *Server) maxSizeWriter(req *http.Request, w io.Writer) (io.Writer, int, if err != nil { return nil, http.StatusLengthRequired, err } - if currentSize+contentLen > s.MaxRepoSize { + if currentSize+contentLen > m.maxRepoSize { err := fmt.Errorf("incoming blob (%d bytes) would exceed maximum size of repository (%d bytes)", - contentLen, s.MaxRepoSize) + contentLen, m.maxRepoSize) return nil, http.StatusRequestEntityTooLarge, err } } // since we can't always trust content-length, we will wrap the writer // in a custom writer that enforces the size limit during writes - return maxSizeWriter{Writer: w, server: s}, 0, nil + return maxSizeWriter{Writer: w, m: m}, 0, nil } // repoSpaceRemaining returns how much space is available in the repo // according to s.MaxRepoSize. s.repoSize must already be set. // If there is no limit, -1 is returned. -func (s *Server) repoSpaceRemaining() int64 { - if s.MaxRepoSize == 0 { +func (m *Manager) repoSpaceRemaining() int64 { + if m.maxRepoSize == 0 { return -1 } - maxSize := s.MaxRepoSize - currentSize := atomic.LoadInt64(&s.repoSize) + maxSize := m.maxRepoSize + currentSize := atomic.LoadInt64(&m.repoSize) return maxSize - currentSize } // incrementRepoSpaceUsage increments the current repo size (which // must already be initialized). -func (s *Server) incrementRepoSpaceUsage(by int64) { - atomic.AddInt64(&s.repoSize, by) +func (m *Manager) incrementRepoSpaceUsage(by int64) { + atomic.AddInt64(&m.repoSize, by) +} + +// tallySize counts the size of the contents of path. +func tallySize(path string) (int64, error) { + if path == "" { + path = "." + } + var size int64 + err := filepath.Walk(path, func(path string, info os.FileInfo, err error) error { + if err != nil { + return err + } + size += info.Size() + return nil + }) + return size, err }