services/keepproxy/keepproxy.go

   1 package main
   2
   3 import (
   4         "errors"
   5         "flag"
   6         "fmt"
   7         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
   8         "git.curoverse.com/arvados.git/sdk/go/keepclient"
   9         "github.com/gorilla/mux"
  10         "io"
  11         "io/ioutil"
  12         "log"
  13         "net"
  14         "net/http"
  15         "os"
  16         "os/signal"
  17         "reflect"
  18         "sync"
  19         "syscall"
  20         "time"
  21 )
  22
  23 // Default TCP address on which to listen for requests.
  24 // Initialized by the -listen flag.
  25 const DEFAULT_ADDR = ":25107"
  26
  27 var listener net.Listener
  28
  29 func main() {
  30         var (
  31                 listen           string
  32                 no_get           bool
  33                 no_put           bool
  34                 default_replicas int
  35                 timeout          int64
  36                 pidfile          string
  37         )
  38
  39         flagset := flag.NewFlagSet("default", flag.ExitOnError)
  40
  41         flagset.StringVar(
  42                 &listen,
  43                 "listen",
  44                 DEFAULT_ADDR,
  45                 "Interface on which to listen for requests, in the format "+
  46                         "ipaddr:port. e.g. -listen=10.0.1.24:8000. Use -listen=:port "+
  47                         "to listen on all network interfaces.")
  48
  49         flagset.BoolVar(
  50                 &no_get,
  51                 "no-get",
  52                 false,
  53                 "If set, disable GET operations")
  54
  55         flagset.BoolVar(
  56                 &no_put,
  57                 "no-put",
  58                 false,
  59                 "If set, disable PUT operations")
  60
  61         flagset.IntVar(
  62                 &default_replicas,
  63                 "default-replicas",
  64                 2,
  65                 "Default number of replicas to write if not specified by the client.")
  66
  67         flagset.Int64Var(
  68                 &timeout,
  69                 "timeout",
  70                 15,
  71                 "Timeout on requests to internal Keep services (default 15 seconds)")
  72
  73         flagset.StringVar(
  74                 &pidfile,
  75                 "pid",
  76                 "",
  77                 "Path to write pid file")
  78
  79         flagset.Parse(os.Args[1:])
  80
  81         arv, err := arvadosclient.MakeArvadosClient()
  82         if err != nil {
  83                 log.Fatalf("Error setting up arvados client %s", err.Error())
  84         }
  85
  86         kc, err := keepclient.MakeKeepClient(&arv)
  87         if err != nil {
  88                 log.Fatalf("Error setting up keep client %s", err.Error())
  89         }
  90
  91         if pidfile != "" {
  92                 f, err := os.Create(pidfile)
  93                 if err != nil {
  94                         log.Fatalf("Error writing pid file (%s): %s", pidfile, err.Error())
  95                 }
  96                 fmt.Fprint(f, os.Getpid())
  97                 f.Close()
  98                 defer os.Remove(pidfile)
  99         }
 100
 101         kc.Want_replicas = default_replicas
 102
 103         kc.Client.Timeout = time.Duration(timeout) * time.Second
 104
 105         listener, err = net.Listen("tcp", listen)
 106         if err != nil {
 107                 log.Fatalf("Could not listen on %v", listen)
 108         }
 109
 110         go RefreshServicesList(kc)
 111
 112         // Shut down the server gracefully (by closing the listener)
 113         // if SIGTERM is received.
 114         term := make(chan os.Signal, 1)
 115         go func(sig <-chan os.Signal) {
 116                 s := <-sig
 117                 log.Println("caught signal:", s)
 118                 listener.Close()
 119         }(term)
 120         signal.Notify(term, syscall.SIGTERM)
 121         signal.Notify(term, syscall.SIGINT)
 122
 123         log.Printf("Arvados Keep proxy started listening on %v", listener.Addr())
 124
 125         // Start listening for requests.
 126         http.Serve(listener, MakeRESTRouter(!no_get, !no_put, kc))
 127
 128         log.Println("shutting down")
 129 }
 130
 131 type ApiTokenCache struct {
 132         tokens     map[string]int64
 133         lock       sync.Mutex
 134         expireTime int64
 135 }
 136
 137 // Refresh the keep service list every five minutes.
 138 func RefreshServicesList(kc *keepclient.KeepClient) {
 139         var previousRoots = []map[string]string{}
 140         var delay time.Duration = 0
 141         for {
 142                 time.Sleep(delay * time.Second)
 143                 delay = 300
 144                 if err := kc.DiscoverKeepServers(); err != nil {
 145                         log.Println("Error retrieving services list:", err)
 146                         delay = 3
 147                         continue
 148                 }
 149                 newRoots := []map[string]string{kc.LocalRoots(), kc.GatewayRoots()}
 150                 if !reflect.DeepEqual(previousRoots, newRoots) {
 151                         log.Printf("Updated services list: locals %v gateways %v", newRoots[0], newRoots[1])
 152                 }
 153                 if len(newRoots[0]) == 0 {
 154                         log.Print("WARNING: No local services. Retrying in 3 seconds.")
 155                         delay = 3
 156                 }
 157                 previousRoots = newRoots
 158         }
 159 }
 160
 161 // Cache the token and set an expire time.  If we already have an expire time
 162 // on the token, it is not updated.
 163 func (this *ApiTokenCache) RememberToken(token string) {
 164         this.lock.Lock()
 165         defer this.lock.Unlock()
 166
 167         now := time.Now().Unix()
 168         if this.tokens[token] == 0 {
 169                 this.tokens[token] = now + this.expireTime
 170         }
 171 }
 172
 173 // Check if the cached token is known and still believed to be valid.
 174 func (this *ApiTokenCache) RecallToken(token string) bool {
 175         this.lock.Lock()
 176         defer this.lock.Unlock()
 177
 178         now := time.Now().Unix()
 179         if this.tokens[token] == 0 {
 180                 // Unknown token
 181                 return false
 182         } else if now < this.tokens[token] {
 183                 // Token is known and still valid
 184                 return true
 185         } else {
 186                 // Token is expired
 187                 this.tokens[token] = 0
 188                 return false
 189         }
 190 }
 191
 192 func GetRemoteAddress(req *http.Request) string {
 193         if realip := req.Header.Get("X-Real-IP"); realip != "" {
 194                 if forwarded := req.Header.Get("X-Forwarded-For"); forwarded != realip {
 195                         return fmt.Sprintf("%s (X-Forwarded-For %s)", realip, forwarded)
 196                 } else {
 197                         return realip
 198                 }
 199         }
 200         return req.RemoteAddr
 201 }
 202
 203 func CheckAuthorizationHeader(kc keepclient.KeepClient, cache *ApiTokenCache, req *http.Request) (pass bool, tok string) {
 204         var auth string
 205         if auth = req.Header.Get("Authorization"); auth == "" {
 206                 return false, ""
 207         }
 208
 209         _, err := fmt.Sscanf(auth, "OAuth2 %s", &tok)
 210         if err != nil {
 211                 // Scanning error
 212                 return false, ""
 213         }
 214
 215         if cache.RecallToken(tok) {
 216                 // Valid in the cache, short circut
 217                 return true, tok
 218         }
 219
 220         arv := *kc.Arvados
 221         arv.ApiToken = tok
 222         if err := arv.Call("HEAD", "users", "", "current", nil, nil); err != nil {
 223                 log.Printf("%s: CheckAuthorizationHeader error: %v", GetRemoteAddress(req), err)
 224                 return false, ""
 225         }
 226
 227         // Success!  Update cache
 228         cache.RememberToken(tok)
 229
 230         return true, tok
 231 }
 232
 233 type GetBlockHandler struct {
 234         *keepclient.KeepClient
 235         *ApiTokenCache
 236 }
 237
 238 type PutBlockHandler struct {
 239         *keepclient.KeepClient
 240         *ApiTokenCache
 241 }
 242
 243 type InvalidPathHandler struct{}
 244
 245 type OptionsHandler struct{}
 246
 247 // MakeRESTRouter
 248 //     Returns a mux.Router that passes GET and PUT requests to the
 249 //     appropriate handlers.
 250 //
 251 func MakeRESTRouter(
 252         enable_get bool,
 253         enable_put bool,
 254         kc *keepclient.KeepClient) *mux.Router {
 255
 256         t := &ApiTokenCache{tokens: make(map[string]int64), expireTime: 300}
 257
 258         rest := mux.NewRouter()
 259
 260         if enable_get {
 261                 rest.Handle(`/{locator:[0-9a-f]{32}\+.*}`,
 262                         GetBlockHandler{kc, t}).Methods("GET", "HEAD")
 263                 rest.Handle(`/{locator:[0-9a-f]{32}}`, GetBlockHandler{kc, t}).Methods("GET", "HEAD")
 264         }
 265
 266         if enable_put {
 267                 rest.Handle(`/{locator:[0-9a-f]{32}\+.*}`, PutBlockHandler{kc, t}).Methods("PUT")
 268                 rest.Handle(`/{locator:[0-9a-f]{32}}`, PutBlockHandler{kc, t}).Methods("PUT")
 269                 rest.Handle(`/`, PutBlockHandler{kc, t}).Methods("POST")
 270                 rest.Handle(`/{any}`, OptionsHandler{}).Methods("OPTIONS")
 271                 rest.Handle(`/`, OptionsHandler{}).Methods("OPTIONS")
 272         }
 273
 274         rest.NotFoundHandler = InvalidPathHandler{}
 275
 276         return rest
 277 }
 278
 279 func SetCorsHeaders(resp http.ResponseWriter) {
 280         resp.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, POST, PUT, OPTIONS")
 281         resp.Header().Set("Access-Control-Allow-Origin", "*")
 282         resp.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Length, Content-Type, X-Keep-Desired-Replicas")
 283         resp.Header().Set("Access-Control-Max-Age", "86486400")
 284 }
 285
 286 func (this InvalidPathHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 287         log.Printf("%s: %s %s unroutable", GetRemoteAddress(req), req.Method, req.URL.Path)
 288         http.Error(resp, "Bad request", http.StatusBadRequest)
 289 }
 290
 291 func (this OptionsHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 292         log.Printf("%s: %s %s", GetRemoteAddress(req), req.Method, req.URL.Path)
 293         SetCorsHeaders(resp)
 294 }
 295
 296 var BadAuthorizationHeader = errors.New("Missing or invalid Authorization header")
 297 var ContentLengthMismatch = errors.New("Actual length != expected content length")
 298 var MethodNotSupported = errors.New("Method not supported")
 299
 300 func (this GetBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 301         SetCorsHeaders(resp)
 302
 303         locator := mux.Vars(req)["locator"]
 304         var err error
 305         var status int
 306         var expectLength, responseLength int64
 307         var proxiedURI = "-"
 308
 309         defer func() {
 310                 log.Println(GetRemoteAddress(req), req.Method, req.URL.Path, status, expectLength, responseLength, proxiedURI, err)
 311                 if status != http.StatusOK {
 312                         http.Error(resp, err.Error(), status)
 313                 }
 314         }()
 315
 316         kc := *this.KeepClient
 317
 318         var pass bool
 319         var tok string
 320         if pass, tok = CheckAuthorizationHeader(kc, this.ApiTokenCache, req); !pass {
 321                 status, err = http.StatusForbidden, BadAuthorizationHeader
 322                 return
 323         }
 324
 325         // Copy ArvadosClient struct and use the client's API token
 326         arvclient := *kc.Arvados
 327         arvclient.ApiToken = tok
 328         kc.Arvados = &arvclient
 329
 330         var reader io.ReadCloser
 331
 332         switch req.Method {
 333         case "HEAD":
 334                 expectLength, proxiedURI, err = kc.Ask(locator)
 335         case "GET":
 336                 reader, expectLength, proxiedURI, err = kc.Get(locator)
 337                 if reader != nil {
 338                         defer reader.Close()
 339                 }
 340         default:
 341                 status, err = http.StatusNotImplemented, MethodNotSupported
 342                 return
 343         }
 344
 345         if expectLength == -1 {
 346                 log.Println("Warning:", GetRemoteAddress(req), req.Method, proxiedURI, "Content-Length not provided")
 347         }
 348
 349         switch err {
 350         case nil:
 351                 status = http.StatusOK
 352                 resp.Header().Set("Content-Length", fmt.Sprint(expectLength))
 353                 switch req.Method {
 354                 case "HEAD":
 355                         responseLength = 0
 356                 case "GET":
 357                         responseLength, err = io.Copy(resp, reader)
 358                         if err == nil && expectLength > -1 && responseLength != expectLength {
 359                                 err = ContentLengthMismatch
 360                         }
 361                 }
 362         case keepclient.BlockNotFound:
 363                 status = http.StatusNotFound
 364         default:
 365                 status = http.StatusBadGateway
 366         }
 367 }
 368
 369 var LengthRequiredError = errors.New(http.StatusText(http.StatusLengthRequired))
 370 var LengthMismatchError = errors.New("Locator size hint does not match Content-Length header")
 371
 372 func (this PutBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 373         SetCorsHeaders(resp)
 374
 375         kc := *this.KeepClient
 376         var err error
 377         var expectLength int64 = -1
 378         var status = http.StatusInternalServerError
 379         var wroteReplicas int
 380         var locatorOut string = "-"
 381
 382         defer func() {
 383                 log.Println(GetRemoteAddress(req), req.Method, req.URL.Path, status, expectLength, kc.Want_replicas, wroteReplicas, locatorOut, err)
 384                 if status != http.StatusOK {
 385                         http.Error(resp, err.Error(), status)
 386                 }
 387         }()
 388
 389         locatorIn := mux.Vars(req)["locator"]
 390
 391         if req.Header.Get("Content-Length") != "" {
 392                 _, err := fmt.Sscanf(req.Header.Get("Content-Length"), "%d", &expectLength)
 393                 if err != nil {
 394                         resp.Header().Set("Content-Length", fmt.Sprintf("%d", expectLength))
 395                 }
 396
 397         }
 398
 399         if expectLength < 0 {
 400                 err = LengthRequiredError
 401                 status = http.StatusLengthRequired
 402                 return
 403         }
 404
 405         if locatorIn != "" {
 406                 var loc *keepclient.Locator
 407                 if loc, err = keepclient.MakeLocator(locatorIn); err != nil {
 408                         status = http.StatusBadRequest
 409                         return
 410                 } else if loc.Size > 0 && int64(loc.Size) != expectLength {
 411                         err = LengthMismatchError
 412                         status = http.StatusBadRequest
 413                         return
 414                 }
 415         }
 416
 417         var pass bool
 418         var tok string
 419         if pass, tok = CheckAuthorizationHeader(kc, this.ApiTokenCache, req); !pass {
 420                 err = BadAuthorizationHeader
 421                 status = http.StatusForbidden
 422                 return
 423         }
 424
 425         // Copy ArvadosClient struct and use the client's API token
 426         arvclient := *kc.Arvados
 427         arvclient.ApiToken = tok
 428         kc.Arvados = &arvclient
 429
 430         // Check if the client specified the number of replicas
 431         if req.Header.Get("X-Keep-Desired-Replicas") != "" {
 432                 var r int
 433                 _, err := fmt.Sscanf(req.Header.Get(keepclient.X_Keep_Desired_Replicas), "%d", &r)
 434                 if err != nil {
 435                         kc.Want_replicas = r
 436                 }
 437         }
 438
 439         // Now try to put the block through
 440         if locatorIn == "" {
 441                 if bytes, err := ioutil.ReadAll(req.Body); err != nil {
 442                         err = errors.New(fmt.Sprintf("Error reading request body: %s", err))
 443                         status = http.StatusInternalServerError
 444                         return
 445                 } else {
 446                         locatorOut, wroteReplicas, err = kc.PutB(bytes)
 447                 }
 448         } else {
 449                 locatorOut, wroteReplicas, err = kc.PutHR(locatorIn, req.Body, expectLength)
 450         }
 451
 452         // Tell the client how many successful PUTs we accomplished
 453         resp.Header().Set(keepclient.X_Keep_Replicas_Stored, fmt.Sprintf("%d", wroteReplicas))
 454
 455         switch err {
 456         case nil:
 457                 status = http.StatusOK
 458                 _, err = io.WriteString(resp, locatorOut)
 459
 460         case keepclient.OversizeBlockError:
 461                 // Too much data
 462                 status = http.StatusRequestEntityTooLarge
 463
 464         case keepclient.InsufficientReplicasError:
 465                 if wroteReplicas > 0 {
 466                         // At least one write is considered success.  The
 467                         // client can decide if getting less than the number of
 468                         // replications it asked for is a fatal error.
 469                         status = http.StatusOK
 470                         _, err = io.WriteString(resp, locatorOut)
 471                 } else {
 472                         status = http.StatusServiceUnavailable
 473                 }
 474
 475         default:
 476                 status = http.StatusBadGateway
 477         }
 478 }