go-hep.org/x/hep@v0.38.1/groot/rsrv/server.go (about)

     1  // Copyright ©2018 The go-hep Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package rsrv
     6  
     7  import (
     8  	"fmt"
     9  	"log"
    10  	"net/http"
    11  	"path/filepath"
    12  	"sync"
    13  	"time"
    14  
    15  	uuid "github.com/hashicorp/go-uuid"
    16  )
    17  
    18  const (
    19  	cookieName = "GROOT_SRV"
    20  )
    21  
    22  // Server serves and manages ROOT files.
    23  type Server struct {
    24  	quit chan int
    25  
    26  	mu       sync.RWMutex
    27  	cookies  map[string]*http.Cookie
    28  	sessions map[string]*DB
    29  
    30  	dir string
    31  }
    32  
    33  // New creates a new server.
    34  func New(dir string) *Server {
    35  	srv := &Server{
    36  		quit:     make(chan int),
    37  		cookies:  make(map[string]*http.Cookie),
    38  		sessions: make(map[string]*DB),
    39  		dir:      dir,
    40  	}
    41  
    42  	go srv.run()
    43  	return srv
    44  }
    45  
    46  // Shutdown shuts the server down.
    47  func (srv *Server) Shutdown() {
    48  	srv.mu.Lock()
    49  	defer srv.mu.Unlock()
    50  	for name := range srv.cookies {
    51  		srv.sessions[name].Close()
    52  		delete(srv.sessions, name)
    53  		delete(srv.cookies, name)
    54  	}
    55  	close(srv.quit)
    56  }
    57  
    58  func (srv *Server) run() {
    59  	ticker := time.NewTicker(5 * time.Minute)
    60  	defer ticker.Stop()
    61  	srv.gc()
    62  	for {
    63  		select {
    64  		case <-ticker.C:
    65  			srv.gc()
    66  		case <-srv.quit:
    67  			return
    68  		}
    69  	}
    70  }
    71  
    72  func (srv *Server) gc() {
    73  	srv.mu.Lock()
    74  	defer srv.mu.Unlock()
    75  	for name, cookie := range srv.cookies {
    76  		now := time.Now()
    77  		if now.After(cookie.Expires) {
    78  			srv.sessions[name].Close()
    79  			delete(srv.sessions, name)
    80  			delete(srv.cookies, name)
    81  			cookie.MaxAge = -1
    82  		}
    83  	}
    84  }
    85  
    86  func (srv *Server) wrap(fn func(w http.ResponseWriter, r *http.Request) error) http.HandlerFunc {
    87  	return func(w http.ResponseWriter, r *http.Request) {
    88  		err := srv.setCookie(w, r)
    89  		if err != nil {
    90  			log.Printf("error retrieving cookie: %v\n", err)
    91  			http.Error(w, err.Error(), http.StatusInternalServerError)
    92  			return
    93  		}
    94  
    95  		if err := fn(w, r); err != nil {
    96  			log.Printf("error %q: %v\n", r.URL.Path, err.Error())
    97  			http.Error(w, err.Error(), http.StatusInternalServerError)
    98  		}
    99  	}
   100  }
   101  
   102  func (srv *Server) setCookie(w http.ResponseWriter, r *http.Request) error {
   103  	srv.mu.Lock()
   104  	defer srv.mu.Unlock()
   105  	cookie, err := r.Cookie(cookieName)
   106  	if err != nil && err != http.ErrNoCookie {
   107  		return err
   108  	}
   109  
   110  	if cookie != nil {
   111  		if v, ok := srv.sessions[cookie.Value]; v == nil || !ok {
   112  			srv.sessions[cookie.Value] = NewDB(filepath.Join(srv.dir, cookie.Value))
   113  			srv.cookies[cookie.Value] = cookie
   114  		}
   115  		return nil
   116  	}
   117  
   118  	v, err := uuid.GenerateUUID()
   119  	if err != nil {
   120  		return fmt.Errorf("could not generate UUID: %w", err)
   121  	}
   122  
   123  	cookie = &http.Cookie{
   124  		Name:    cookieName,
   125  		Value:   v,
   126  		Expires: time.Now().Add(24 * time.Hour),
   127  	}
   128  	srv.sessions[cookie.Value] = NewDB(filepath.Join(srv.dir, cookie.Value))
   129  	srv.cookies[cookie.Value] = cookie
   130  	http.SetCookie(w, cookie)
   131  	return nil
   132  }
   133  
   134  func (srv *Server) cookie(r *http.Request) (*http.Cookie, error) {
   135  	srv.mu.RLock()
   136  	defer srv.mu.RUnlock()
   137  	cookie, err := r.Cookie(cookieName)
   138  	if err != nil {
   139  		return nil, err
   140  	}
   141  
   142  	if cookie == nil {
   143  		return nil, http.ErrNoCookie
   144  	}
   145  	return srv.cookies[cookie.Value], nil
   146  }
   147  
   148  func (srv *Server) db(r *http.Request) (*DB, error) {
   149  	srv.mu.RLock()
   150  	defer srv.mu.RUnlock()
   151  	cookie, err := srv.cookie(r)
   152  	if err != nil {
   153  		return nil, err
   154  	}
   155  	if cookie == nil {
   156  		return nil, http.ErrNoCookie
   157  	}
   158  	return srv.sessions[cookie.Value], nil
   159  }
   160  
   161  // DB returns the underlying data base of files associated with the user
   162  // identified by their cookie.
   163  func (srv *Server) DB(r *http.Request) (*DB, error) {
   164  	return srv.db(r)
   165  }