services/keep/src/arvados.org/keepproxy/keepproxy.go

   1 package main
   2
   3 import (
   4         "arvados.org/keepclient"
   5         "flag"
   6         "fmt"
   7         "github.com/gorilla/mux"
   8         "io"
   9         "log"
  10         "net"
  11         "net/http"
  12         "os"
  13         "sync"
  14         "time"
  15 )
  16
  17 // Default TCP address on which to listen for requests.
  18 // Initialized by the -listen flag.
  19 const DEFAULT_ADDR = ":25107"
  20
  21 var listener net.Listener
  22
  23 func main() {
  24         var (
  25                 listen           string
  26                 no_get           bool
  27                 no_put           bool
  28                 default_replicas int
  29                 pidfile          string
  30         )
  31
  32         flag.StringVar(
  33                 &listen,
  34                 "listen",
  35                 DEFAULT_ADDR,
  36                 "Interface on which to listen for requests, in the format "+
  37                         "ipaddr:port. e.g. -listen=10.0.1.24:8000. Use -listen=:port "+
  38                         "to listen on all network interfaces.")
  39
  40         flag.BoolVar(
  41                 &no_get,
  42                 "no-get",
  43                 false,
  44                 "If set, disable GET operations")
  45
  46         flag.BoolVar(
  47                 &no_get,
  48                 "no-put",
  49                 false,
  50                 "If set, disable PUT operations")
  51
  52         flag.IntVar(
  53                 &default_replicas,
  54                 "default-replicas",
  55                 2,
  56                 "Default number of replicas to write if not specified by the client.")
  57
  58         flag.StringVar(
  59                 &pidfile,
  60                 "pid",
  61                 "",
  62                 "Path to write pid file")
  63
  64         flag.Parse()
  65
  66         kc, err := keepclient.MakeKeepClient()
  67         if err != nil {
  68                 log.Fatalf("Error setting up keep client %s", err.Error())
  69         }
  70
  71         if pidfile != "" {
  72                 f, err := os.Create(pidfile)
  73                 if err == nil {
  74                         fmt.Fprint(f, os.Getpid())
  75                         f.Close()
  76                 } else {
  77                         log.Printf("Error writing pid file (%s): %s", pidfile, err.Error())
  78                 }
  79         }
  80
  81         kc.Want_replicas = default_replicas
  82
  83         listener, err = net.Listen("tcp", listen)
  84         if err != nil {
  85                 log.Fatalf("Could not listen on %v", listen)
  86         }
  87
  88         go RefreshServicesList(&kc)
  89
  90         log.Printf("Arvados Keep proxy started listening on %v with server list %v", listener.Addr(), kc.ServiceRoots())
  91
  92         // Start listening for requests.
  93         http.Serve(listener, MakeRESTRouter(!no_get, !no_put, &kc))
  94 }
  95
  96 type ApiTokenCache struct {
  97         tokens     map[string]int64
  98         lock       sync.Mutex
  99         expireTime int64
 100 }
 101
 102 // Refresh the keep service list every five minutes.
 103 func RefreshServicesList(kc *keepclient.KeepClient) {
 104         for {
 105                 time.Sleep(300 * time.Second)
 106                 oldservices := kc.ServiceRoots()
 107                 kc.DiscoverKeepServers()
 108                 newservices := kc.ServiceRoots()
 109                 s1 := fmt.Sprint(oldservices)
 110                 s2 := fmt.Sprint(newservices)
 111                 if s1 != s2 {
 112                         log.Printf("Updated server list to %v", s2)
 113                 }
 114         }
 115 }
 116
 117 // Cache the token and set an expire time.  If we already have an expire time
 118 // on the token, it is not updated.
 119 func (this *ApiTokenCache) RememberToken(token string) {
 120         this.lock.Lock()
 121         defer this.lock.Unlock()
 122
 123         now := time.Now().Unix()
 124         if this.tokens[token] == 0 {
 125                 this.tokens[token] = now + this.expireTime
 126         }
 127 }
 128
 129 // Check if the cached token is known and still believed to be valid.
 130 func (this *ApiTokenCache) RecallToken(token string) bool {
 131         this.lock.Lock()
 132         defer this.lock.Unlock()
 133
 134         now := time.Now().Unix()
 135         if this.tokens[token] == 0 {
 136                 // Unknown token
 137                 return false
 138         } else if now < this.tokens[token] {
 139                 // Token is known and still valid
 140                 return true
 141         } else {
 142                 // Token is expired
 143                 this.tokens[token] = 0
 144                 return false
 145         }
 146 }
 147
 148 func GetRemoteAddress(req *http.Request) string {
 149         if realip := req.Header.Get("X-Real-IP"); realip != "" {
 150                 if forwarded := req.Header.Get("X-Forwarded-For"); forwarded != realip {
 151                         return fmt.Sprintf("%s (X-Forwarded-For %s)", realip, forwarded)
 152                 } else {
 153                         return realip
 154                 }
 155         }
 156         return req.RemoteAddr
 157 }
 158
 159 func CheckAuthorizationHeader(kc keepclient.KeepClient, cache *ApiTokenCache, req *http.Request) bool {
 160         var auth string
 161         if auth = req.Header.Get("Authorization"); auth == "" {
 162                 return false
 163         }
 164
 165         var tok string
 166         _, err := fmt.Sscanf(auth, "OAuth2 %s", &tok)
 167         if err != nil {
 168                 // Scanning error
 169                 return false
 170         }
 171
 172         if cache.RecallToken(tok) {
 173                 // Valid in the cache, short circut
 174                 return true
 175         }
 176
 177         var usersreq *http.Request
 178
 179         if usersreq, err = http.NewRequest("HEAD", fmt.Sprintf("https://%s/arvados/v1/users/current", kc.ApiServer), nil); err != nil {
 180                 // Can't construct the request
 181                 log.Printf("%s: CheckAuthorizationHeader error: %v", GetRemoteAddress(req), err)
 182                 return false
 183         }
 184
 185         // Add api token header
 186         usersreq.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", tok))
 187
 188         // Actually make the request
 189         var resp *http.Response
 190         if resp, err = kc.Client.Do(usersreq); err != nil {
 191                 // Something else failed
 192                 log.Printf("%s: CheckAuthorizationHeader error connecting to API server: %v", GetRemoteAddress(req), err.Error())
 193                 return false
 194         }
 195
 196         if resp.StatusCode != http.StatusOK {
 197                 // Bad status
 198                 log.Printf("%s: CheckAuthorizationHeader API server responded: %v", GetRemoteAddress(req), resp.Status)
 199                 return false
 200         }
 201
 202         // Success!  Update cache
 203         cache.RememberToken(tok)
 204
 205         return true
 206 }
 207
 208 type GetBlockHandler struct {
 209         *keepclient.KeepClient
 210         *ApiTokenCache
 211 }
 212
 213 type PutBlockHandler struct {
 214         *keepclient.KeepClient
 215         *ApiTokenCache
 216 }
 217
 218 type InvalidPathHandler struct{}
 219
 220 // MakeRESTRouter
 221 //     Returns a mux.Router that passes GET and PUT requests to the
 222 //     appropriate handlers.
 223 //
 224 func MakeRESTRouter(
 225         enable_get bool,
 226         enable_put bool,
 227         kc *keepclient.KeepClient) *mux.Router {
 228
 229         t := &ApiTokenCache{tokens: make(map[string]int64), expireTime: 300}
 230
 231         rest := mux.NewRouter()
 232         gh := rest.Handle(`/{hash:[0-9a-f]{32}}`, GetBlockHandler{kc, t})
 233         ghsig := rest.Handle(
 234                 `/{hash:[0-9a-f]{32}}+A{signature:[0-9a-f]+}@{timestamp:[0-9a-f]+}`,
 235                 GetBlockHandler{kc, t})
 236         ph := rest.Handle(`/{hash:[0-9a-f]{32}}`, PutBlockHandler{kc, t})
 237
 238         if enable_get {
 239                 gh.Methods("GET", "HEAD")
 240                 ghsig.Methods("GET", "HEAD")
 241         }
 242
 243         if enable_put {
 244                 ph.Methods("PUT")
 245         }
 246
 247         rest.NotFoundHandler = InvalidPathHandler{}
 248
 249         return rest
 250 }
 251
 252 func (this InvalidPathHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 253         log.Printf("%s: %s %s unroutable", GetRemoteAddress(req), req.Method, req.URL.Path)
 254         http.Error(resp, "Bad request", http.StatusBadRequest)
 255 }
 256
 257 func (this GetBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 258
 259         kc := *this.KeepClient
 260
 261         hash := mux.Vars(req)["hash"]
 262         signature := mux.Vars(req)["signature"]
 263         timestamp := mux.Vars(req)["timestamp"]
 264
 265         log.Printf("%s: %s %s", GetRemoteAddress(req), req.Method, hash)
 266
 267         if !CheckAuthorizationHeader(kc, this.ApiTokenCache, req) {
 268                 http.Error(resp, "Missing or invalid Authorization header", http.StatusForbidden)
 269                 return
 270         }
 271
 272         var reader io.ReadCloser
 273         var err error
 274         var blocklen int64
 275
 276         if req.Method == "GET" {
 277                 reader, blocklen, _, err = kc.AuthorizedGet(hash, signature, timestamp)
 278                 defer reader.Close()
 279         } else if req.Method == "HEAD" {
 280                 blocklen, _, err = kc.AuthorizedAsk(hash, signature, timestamp)
 281         }
 282
 283         resp.Header().Set("Content-Length", fmt.Sprint(blocklen))
 284
 285         switch err {
 286         case nil:
 287                 if reader != nil {
 288                         n, err2 := io.Copy(resp, reader)
 289                         if n != blocklen {
 290                                 log.Printf("%s: %s %s mismatched return %v with Content-Length %v error", GetRemoteAddress(req), req.Method, hash, n, blocklen, err.Error())
 291                         } else if err2 == nil {
 292                                 log.Printf("%s: %s %s success returned %v bytes", GetRemoteAddress(req), req.Method, hash, n)
 293                         } else {
 294                                 log.Printf("%s: %s %s returned %v bytes error %v", GetRemoteAddress(req), req.Method, hash, n, err.Error())
 295                         }
 296                 } else {
 297                         log.Printf("%s: %s %s success", GetRemoteAddress(req), req.Method, hash)
 298                 }
 299         case keepclient.BlockNotFound:
 300                 http.Error(resp, "Not found", http.StatusNotFound)
 301         default:
 302                 http.Error(resp, err.Error(), http.StatusBadGateway)
 303         }
 304
 305         if err != nil {
 306                 log.Printf("%s: %s %s error %s", GetRemoteAddress(req), req.Method, hash, err.Error())
 307         }
 308 }
 309
 310 func (this PutBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 311
 312         kc := *this.KeepClient
 313
 314         hash := mux.Vars(req)["hash"]
 315
 316         var contentLength int64 = -1
 317         if req.Header.Get("Content-Length") != "" {
 318                 _, err := fmt.Sscanf(req.Header.Get("Content-Length"), "%d", &contentLength)
 319                 if err != nil {
 320                         resp.Header().Set("Content-Length", fmt.Sprintf("%d", contentLength))
 321                 }
 322
 323         }
 324
 325         log.Printf("%s: %s %s Content-Length %v", GetRemoteAddress(req), req.Method, hash, contentLength)
 326
 327         if contentLength < 1 {
 328                 http.Error(resp, "Must include Content-Length header", http.StatusLengthRequired)
 329                 return
 330         }
 331
 332         if !CheckAuthorizationHeader(kc, this.ApiTokenCache, req) {
 333                 http.Error(resp, "Missing or invalid Authorization header", http.StatusForbidden)
 334                 return
 335         }
 336
 337         // Check if the client specified the number of replicas
 338         if req.Header.Get("X-Keep-Desired-Replicas") != "" {
 339                 var r int
 340                 _, err := fmt.Sscanf(req.Header.Get(keepclient.X_Keep_Desired_Replicas), "%d", &r)
 341                 if err != nil {
 342                         kc.Want_replicas = r
 343                 }
 344         }
 345
 346         // Now try to put the block through
 347         replicas, err := kc.PutHR(hash, req.Body, contentLength)
 348
 349         // Tell the client how many successful PUTs we accomplished
 350         resp.Header().Set(keepclient.X_Keep_Replicas_Stored, fmt.Sprintf("%d", replicas))
 351
 352         switch err {
 353         case nil:
 354                 // Default will return http.StatusOK
 355                 log.Printf("%s: %s %s finished, stored %v replicas (desired %v)", GetRemoteAddress(req), req.Method, hash, replicas, kc.Want_replicas)
 356
 357         case keepclient.OversizeBlockError:
 358                 // Too much data
 359                 http.Error(resp, fmt.Sprintf("Exceeded maximum blocksize %d", keepclient.BLOCKSIZE), http.StatusRequestEntityTooLarge)
 360
 361         case keepclient.InsufficientReplicasError:
 362                 if replicas > 0 {
 363                         // At least one write is considered success.  The
 364                         // client can decide if getting less than the number of
 365                         // replications it asked for is a fatal error.
 366                         // Default will return http.StatusOK
 367                 } else {
 368                         http.Error(resp, "", http.StatusServiceUnavailable)
 369                 }
 370
 371         default:
 372                 http.Error(resp, err.Error(), http.StatusBadGateway)
 373         }
 374
 375         if err != nil {
 376                 log.Printf("%s: %s %s stored %v replicas (desired %v) got error %v", GetRemoteAddress(req), req.Method, hash, replicas, kc.Want_replicas, err.Error())
 377         }
 378
 379 }