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         "regexp"
  18         "sync"
  19         "syscall"
  20         "time"
  21 )
  22
  23 // Default TCP address on which to listen for requests.
  24 // Override with -listen.
  25 const DefaultAddr = ":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("keepproxy", flag.ExitOnError)
  40
  41         flagset.StringVar(
  42                 &listen,
  43                 "listen",
  44                 DefaultAddr,
  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         if os.Getenv("ARVADOS_DEBUG") != "" {
  87                 keepclient.DebugPrintf = log.Printf
  88         }
  89         kc, err := keepclient.MakeKeepClient(&arv)
  90         if err != nil {
  91                 log.Fatalf("Error setting up keep client %s", err.Error())
  92         }
  93
  94         if pidfile != "" {
  95                 f, err := os.Create(pidfile)
  96                 if err != nil {
  97                         log.Fatalf("Error writing pid file (%s): %s", pidfile, err.Error())
  98                 }
  99                 fmt.Fprint(f, os.Getpid())
 100                 f.Close()
 101                 defer os.Remove(pidfile)
 102         }
 103
 104         kc.Want_replicas = default_replicas
 105         kc.Client.Timeout = time.Duration(timeout) * time.Second
 106         go kc.RefreshServices(5*time.Minute, 3*time.Second)
 107
 108         listener, err = net.Listen("tcp", listen)
 109         if err != nil {
 110                 log.Fatalf("Could not listen on %v", listen)
 111         }
 112         log.Printf("Arvados Keep proxy started listening on %v", listener.Addr())
 113
 114         // Shut down the server gracefully (by closing the listener)
 115         // if SIGTERM is received.
 116         term := make(chan os.Signal, 1)
 117         go func(sig <-chan os.Signal) {
 118                 s := <-sig
 119                 log.Println("caught signal:", s)
 120                 listener.Close()
 121         }(term)
 122         signal.Notify(term, syscall.SIGTERM)
 123         signal.Notify(term, syscall.SIGINT)
 124
 125         // Start serving 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 // Cache the token and set an expire time.  If we already have an expire time
 138 // on the token, it is not updated.
 139 func (this *ApiTokenCache) RememberToken(token string) {
 140         this.lock.Lock()
 141         defer this.lock.Unlock()
 142
 143         now := time.Now().Unix()
 144         if this.tokens[token] == 0 {
 145                 this.tokens[token] = now + this.expireTime
 146         }
 147 }
 148
 149 // Check if the cached token is known and still believed to be valid.
 150 func (this *ApiTokenCache) RecallToken(token string) bool {
 151         this.lock.Lock()
 152         defer this.lock.Unlock()
 153
 154         now := time.Now().Unix()
 155         if this.tokens[token] == 0 {
 156                 // Unknown token
 157                 return false
 158         } else if now < this.tokens[token] {
 159                 // Token is known and still valid
 160                 return true
 161         } else {
 162                 // Token is expired
 163                 this.tokens[token] = 0
 164                 return false
 165         }
 166 }
 167
 168 func GetRemoteAddress(req *http.Request) string {
 169         if xff := req.Header.Get("X-Forwarded-For"); xff != "" {
 170                 return xff + "," + req.RemoteAddr
 171         }
 172         return req.RemoteAddr
 173 }
 174
 175 func CheckAuthorizationHeader(kc *keepclient.KeepClient, cache *ApiTokenCache, req *http.Request) (pass bool, tok string) {
 176         var auth string
 177         if auth = req.Header.Get("Authorization"); auth == "" {
 178                 return false, ""
 179         }
 180
 181         _, err := fmt.Sscanf(auth, "OAuth2 %s", &tok)
 182         if err != nil {
 183                 // Scanning error
 184                 return false, ""
 185         }
 186
 187         if cache.RecallToken(tok) {
 188                 // Valid in the cache, short circuit
 189                 return true, tok
 190         }
 191
 192         arv := *kc.Arvados
 193         arv.ApiToken = tok
 194         if err := arv.Call("HEAD", "users", "", "current", nil, nil); err != nil {
 195                 log.Printf("%s: CheckAuthorizationHeader error: %v", GetRemoteAddress(req), err)
 196                 return false, ""
 197         }
 198
 199         // Success!  Update cache
 200         cache.RememberToken(tok)
 201
 202         return true, tok
 203 }
 204
 205 type GetBlockHandler struct {
 206         *keepclient.KeepClient
 207         *ApiTokenCache
 208 }
 209
 210 type PutBlockHandler struct {
 211         *keepclient.KeepClient
 212         *ApiTokenCache
 213 }
 214
 215 type IndexHandler struct {
 216         *keepclient.KeepClient
 217         *ApiTokenCache
 218 }
 219
 220 type InvalidPathHandler struct{}
 221
 222 type OptionsHandler struct{}
 223
 224 // MakeRESTRouter
 225 //     Returns a mux.Router that passes GET and PUT requests to the
 226 //     appropriate handlers.
 227 //
 228 func MakeRESTRouter(
 229         enable_get bool,
 230         enable_put bool,
 231         kc *keepclient.KeepClient) *mux.Router {
 232
 233         t := &ApiTokenCache{tokens: make(map[string]int64), expireTime: 300}
 234
 235         rest := mux.NewRouter()
 236
 237         if enable_get {
 238                 rest.Handle(`/{locator:[0-9a-f]{32}\+.*}`,
 239                         GetBlockHandler{kc, t}).Methods("GET", "HEAD")
 240                 rest.Handle(`/{locator:[0-9a-f]{32}}`, GetBlockHandler{kc, t}).Methods("GET", "HEAD")
 241
 242                 // List all blocks
 243                 rest.Handle(`/index`, IndexHandler{kc, t}).Methods("GET")
 244
 245                 // List blocks whose hash has the given prefix
 246                 rest.Handle(`/index/{prefix:[0-9a-f]{0,32}}`, IndexHandler{kc, t}).Methods("GET")
 247         }
 248
 249         if enable_put {
 250                 rest.Handle(`/{locator:[0-9a-f]{32}\+.*}`, PutBlockHandler{kc, t}).Methods("PUT")
 251                 rest.Handle(`/{locator:[0-9a-f]{32}}`, PutBlockHandler{kc, t}).Methods("PUT")
 252                 rest.Handle(`/`, PutBlockHandler{kc, t}).Methods("POST")
 253                 rest.Handle(`/{any}`, OptionsHandler{}).Methods("OPTIONS")
 254                 rest.Handle(`/`, OptionsHandler{}).Methods("OPTIONS")
 255         }
 256
 257         rest.NotFoundHandler = InvalidPathHandler{}
 258
 259         return rest
 260 }
 261
 262 func SetCorsHeaders(resp http.ResponseWriter) {
 263         resp.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, POST, PUT, OPTIONS")
 264         resp.Header().Set("Access-Control-Allow-Origin", "*")
 265         resp.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Length, Content-Type, X-Keep-Desired-Replicas")
 266         resp.Header().Set("Access-Control-Max-Age", "86486400")
 267 }
 268
 269 func (this InvalidPathHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 270         log.Printf("%s: %s %s unroutable", GetRemoteAddress(req), req.Method, req.URL.Path)
 271         http.Error(resp, "Bad request", http.StatusBadRequest)
 272 }
 273
 274 func (this OptionsHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 275         log.Printf("%s: %s %s", GetRemoteAddress(req), req.Method, req.URL.Path)
 276         SetCorsHeaders(resp)
 277 }
 278
 279 var BadAuthorizationHeader = errors.New("Missing or invalid Authorization header")
 280 var ContentLengthMismatch = errors.New("Actual length != expected content length")
 281 var MethodNotSupported = errors.New("Method not supported")
 282
 283 var removeHint, _ = regexp.Compile("\\+K@[a-z0-9]{5}(\\+|$)")
 284
 285 func (this GetBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 286         SetCorsHeaders(resp)
 287
 288         locator := mux.Vars(req)["locator"]
 289         var err error
 290         var status int
 291         var expectLength, responseLength int64
 292         var proxiedURI = "-"
 293
 294         defer func() {
 295                 log.Println(GetRemoteAddress(req), req.Method, req.URL.Path, status, expectLength, responseLength, proxiedURI, err)
 296                 if status != http.StatusOK {
 297                         http.Error(resp, err.Error(), status)
 298                 }
 299         }()
 300
 301         kc := *this.KeepClient
 302
 303         var pass bool
 304         var tok string
 305         if pass, tok = CheckAuthorizationHeader(&kc, this.ApiTokenCache, req); !pass {
 306                 status, err = http.StatusForbidden, BadAuthorizationHeader
 307                 return
 308         }
 309
 310         // Copy ArvadosClient struct and use the client's API token
 311         arvclient := *kc.Arvados
 312         arvclient.ApiToken = tok
 313         kc.Arvados = &arvclient
 314
 315         var reader io.ReadCloser
 316
 317         locator = removeHint.ReplaceAllString(locator, "$1")
 318
 319         switch req.Method {
 320         case "HEAD":
 321                 expectLength, proxiedURI, err = kc.Ask(locator)
 322         case "GET":
 323                 reader, expectLength, proxiedURI, err = kc.Get(locator)
 324                 if reader != nil {
 325                         defer reader.Close()
 326                 }
 327         default:
 328                 status, err = http.StatusNotImplemented, MethodNotSupported
 329                 return
 330         }
 331
 332         if expectLength == -1 {
 333                 log.Println("Warning:", GetRemoteAddress(req), req.Method, proxiedURI, "Content-Length not provided")
 334         }
 335
 336         switch respErr := err.(type) {
 337         case nil:
 338                 status = http.StatusOK
 339                 resp.Header().Set("Content-Length", fmt.Sprint(expectLength))
 340                 switch req.Method {
 341                 case "HEAD":
 342                         responseLength = 0
 343                 case "GET":
 344                         responseLength, err = io.Copy(resp, reader)
 345                         if err == nil && expectLength > -1 && responseLength != expectLength {
 346                                 err = ContentLengthMismatch
 347                         }
 348                 }
 349         case keepclient.Error:
 350                 if respErr == keepclient.BlockNotFound {
 351                         status = http.StatusNotFound
 352                 } else if respErr.Temporary() {
 353                         status = http.StatusBadGateway
 354                 } else {
 355                         status = 422
 356                 }
 357         default:
 358                 status = http.StatusInternalServerError
 359         }
 360 }
 361
 362 var LengthRequiredError = errors.New(http.StatusText(http.StatusLengthRequired))
 363 var LengthMismatchError = errors.New("Locator size hint does not match Content-Length header")
 364
 365 func (this PutBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 366         SetCorsHeaders(resp)
 367
 368         kc := *this.KeepClient
 369         var err error
 370         var expectLength int64
 371         var status = http.StatusInternalServerError
 372         var wroteReplicas int
 373         var locatorOut string = "-"
 374
 375         defer func() {
 376                 log.Println(GetRemoteAddress(req), req.Method, req.URL.Path, status, expectLength, kc.Want_replicas, wroteReplicas, locatorOut, err)
 377                 if status != http.StatusOK {
 378                         http.Error(resp, err.Error(), status)
 379                 }
 380         }()
 381
 382         locatorIn := mux.Vars(req)["locator"]
 383
 384         _, err = fmt.Sscanf(req.Header.Get("Content-Length"), "%d", &expectLength)
 385         if err != nil || expectLength < 0 {
 386                 err = LengthRequiredError
 387                 status = http.StatusLengthRequired
 388                 return
 389         }
 390
 391         if locatorIn != "" {
 392                 var loc *keepclient.Locator
 393                 if loc, err = keepclient.MakeLocator(locatorIn); err != nil {
 394                         status = http.StatusBadRequest
 395                         return
 396                 } else if loc.Size > 0 && int64(loc.Size) != expectLength {
 397                         err = LengthMismatchError
 398                         status = http.StatusBadRequest
 399                         return
 400                 }
 401         }
 402
 403         var pass bool
 404         var tok string
 405         if pass, tok = CheckAuthorizationHeader(&kc, this.ApiTokenCache, req); !pass {
 406                 err = BadAuthorizationHeader
 407                 status = http.StatusForbidden
 408                 return
 409         }
 410
 411         // Copy ArvadosClient struct and use the client's API token
 412         arvclient := *kc.Arvados
 413         arvclient.ApiToken = tok
 414         kc.Arvados = &arvclient
 415
 416         // Check if the client specified the number of replicas
 417         if req.Header.Get("X-Keep-Desired-Replicas") != "" {
 418                 var r int
 419                 _, err := fmt.Sscanf(req.Header.Get(keepclient.X_Keep_Desired_Replicas), "%d", &r)
 420                 if err == nil {
 421                         kc.Want_replicas = r
 422                 }
 423         }
 424
 425         // Now try to put the block through
 426         if locatorIn == "" {
 427                 if bytes, err := ioutil.ReadAll(req.Body); err != nil {
 428                         err = errors.New(fmt.Sprintf("Error reading request body: %s", err))
 429                         status = http.StatusInternalServerError
 430                         return
 431                 } else {
 432                         locatorOut, wroteReplicas, err = kc.PutB(bytes)
 433                 }
 434         } else {
 435                 locatorOut, wroteReplicas, err = kc.PutHR(locatorIn, req.Body, expectLength)
 436         }
 437
 438         // Tell the client how many successful PUTs we accomplished
 439         resp.Header().Set(keepclient.X_Keep_Replicas_Stored, fmt.Sprintf("%d", wroteReplicas))
 440
 441         switch err {
 442         case nil:
 443                 status = http.StatusOK
 444                 _, err = io.WriteString(resp, locatorOut)
 445
 446         case keepclient.OversizeBlockError:
 447                 // Too much data
 448                 status = http.StatusRequestEntityTooLarge
 449
 450         case keepclient.InsufficientReplicasError:
 451                 if wroteReplicas > 0 {
 452                         // At least one write is considered success.  The
 453                         // client can decide if getting less than the number of
 454                         // replications it asked for is a fatal error.
 455                         status = http.StatusOK
 456                         _, err = io.WriteString(resp, locatorOut)
 457                 } else {
 458                         status = http.StatusServiceUnavailable
 459                 }
 460
 461         default:
 462                 status = http.StatusBadGateway
 463         }
 464 }
 465
 466 // ServeHTTP implementation for IndexHandler
 467 // Supports only GET requests for /index/{prefix:[0-9a-f]{0,32}}
 468 // For each keep server found in LocalRoots:
 469 //   Invokes GetIndex using keepclient
 470 //   Expects "complete" response (terminating with blank new line)
 471 //   Aborts on any errors
 472 // Concatenates responses from all those keep servers and returns
 473 func (handler IndexHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 474         SetCorsHeaders(resp)
 475
 476         prefix := mux.Vars(req)["prefix"]
 477         var err error
 478         var status int
 479
 480         defer func() {
 481                 if status != http.StatusOK {
 482                         http.Error(resp, err.Error(), status)
 483                 }
 484         }()
 485
 486         kc := *handler.KeepClient
 487
 488         ok, token := CheckAuthorizationHeader(&kc, handler.ApiTokenCache, req)
 489         if !ok {
 490                 status, err = http.StatusForbidden, BadAuthorizationHeader
 491                 return
 492         }
 493
 494         // Copy ArvadosClient struct and use the client's API token
 495         arvclient := *kc.Arvados
 496         arvclient.ApiToken = token
 497         kc.Arvados = &arvclient
 498
 499         // Only GET method is supported
 500         if req.Method != "GET" {
 501                 status, err = http.StatusNotImplemented, MethodNotSupported
 502                 return
 503         }
 504
 505         // Get index from all LocalRoots and write to resp
 506         var reader io.Reader
 507         for uuid := range kc.LocalRoots() {
 508                 reader, err = kc.GetIndex(uuid, prefix)
 509                 if err != nil {
 510                         status = http.StatusBadGateway
 511                         return
 512                 }
 513
 514                 _, err = io.Copy(resp, reader)
 515                 if err != nil {
 516                         status = http.StatusBadGateway
 517                         return
 518                 }
 519         }
 520
 521         // Got index from all the keep servers and wrote to resp
 522         status = http.StatusOK
 523         resp.Write([]byte("\n"))
 524 }