services/keepproxy/keepproxy.go

   1 package main
   2
   3 import (
   4         "errors"
   5         "flag"
   6         "fmt"
   7         "io"
   8         "io/ioutil"
   9         "log"
  10         "net"
  11         "net/http"
  12         "os"
  13         "os/signal"
  14         "regexp"
  15         "sync"
  16         "syscall"
  17         "time"
  18
  19         "git.curoverse.com/arvados.git/sdk/go/arvados"
  20         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
  21         "git.curoverse.com/arvados.git/sdk/go/config"
  22         "git.curoverse.com/arvados.git/sdk/go/keepclient"
  23         "github.com/coreos/go-systemd/daemon"
  24         "github.com/ghodss/yaml"
  25         "github.com/gorilla/mux"
  26 )
  27
  28 type Config struct {
  29         Client          arvados.Client
  30         Listen          string
  31         DisableGet      bool
  32         DisablePut      bool
  33         DefaultReplicas int
  34         Timeout         arvados.Duration
  35         PIDFile         string
  36         Debug           bool
  37 }
  38
  39 func DefaultConfig() *Config {
  40         return &Config{
  41                 Listen:  ":25107",
  42                 Timeout: arvados.Duration(15 * time.Second),
  43         }
  44 }
  45
  46 var listener net.Listener
  47
  48 func main() {
  49         cfg := DefaultConfig()
  50
  51         flagset := flag.NewFlagSet("keepproxy", flag.ExitOnError)
  52         flagset.Usage = usage
  53
  54         const deprecated = " (DEPRECATED -- use config file instead)"
  55         flagset.StringVar(&cfg.Listen, "listen", cfg.Listen, "Local port to listen on."+deprecated)
  56         flagset.BoolVar(&cfg.DisableGet, "no-get", cfg.DisableGet, "Disable GET operations."+deprecated)
  57         flagset.BoolVar(&cfg.DisablePut, "no-put", cfg.DisablePut, "Disable PUT operations."+deprecated)
  58         flagset.IntVar(&cfg.DefaultReplicas, "default-replicas", cfg.DefaultReplicas, "Default number of replicas to write if not specified by the client. If 0, use site default."+deprecated)
  59         flagset.StringVar(&cfg.PIDFile, "pid", cfg.PIDFile, "Path to write pid file."+deprecated)
  60         timeoutSeconds := flagset.Int("timeout", int(time.Duration(cfg.Timeout)/time.Second), "Timeout (in seconds) on requests to internal Keep services."+deprecated)
  61
  62         var cfgPath string
  63         const defaultCfgPath = "/etc/arvados/keepproxy/keepproxy.yml"
  64         flagset.StringVar(&cfgPath, "config", defaultCfgPath, "Configuration file `path`")
  65         dumpConfig := flagset.Bool("dump-config", false, "write current configuration to stdout and exit")
  66         flagset.Parse(os.Args[1:])
  67
  68         err := config.LoadFile(cfg, cfgPath)
  69         if err != nil {
  70                 h := os.Getenv("ARVADOS_API_HOST")
  71                 t := os.Getenv("ARVADOS_API_TOKEN")
  72                 if h == "" || t == "" || !os.IsNotExist(err) || cfgPath != defaultCfgPath {
  73                         log.Fatal(err)
  74                 }
  75                 log.Print("DEPRECATED: No config file found, but ARVADOS_API_HOST and ARVADOS_API_TOKEN environment variables are set. Please use a config file instead.")
  76                 cfg.Client.APIHost = h
  77                 cfg.Client.AuthToken = t
  78                 if regexp.MustCompile("^(?i:1|yes|true)$").MatchString(os.Getenv("ARVADOS_API_HOST_INSECURE")) {
  79                         cfg.Client.Insecure = true
  80                 }
  81                 if y, err := yaml.Marshal(cfg); err == nil && !*dumpConfig {
  82                         log.Print("Current configuration:\n", string(y))
  83                 }
  84                 cfg.Timeout = arvados.Duration(time.Duration(*timeoutSeconds) * time.Second)
  85         }
  86
  87         if *dumpConfig {
  88                 y, err := yaml.Marshal(cfg)
  89                 if err != nil {
  90                         log.Fatal(err)
  91                 }
  92                 os.Stdout.Write(y)
  93                 os.Exit(0)
  94         }
  95
  96         arv, err := arvadosclient.New(&cfg.Client)
  97         if err != nil {
  98                 log.Fatalf("Error setting up arvados client %s", err.Error())
  99         }
 100
 101         if cfg.Debug {
 102                 keepclient.DebugPrintf = log.Printf
 103         }
 104         kc, err := keepclient.MakeKeepClient(arv)
 105         if err != nil {
 106                 log.Fatalf("Error setting up keep client %s", err.Error())
 107         }
 108
 109         if cfg.PIDFile != "" {
 110                 f, err := os.Create(cfg.PIDFile)
 111                 if err != nil {
 112                         log.Fatal(err)
 113                 }
 114                 defer f.Close()
 115                 err = syscall.Flock(int(f.Fd()), syscall.LOCK_EX|syscall.LOCK_NB)
 116                 if err != nil {
 117                         log.Fatalf("flock(%s): %s", cfg.PIDFile, err)
 118                 }
 119                 defer os.Remove(cfg.PIDFile)
 120                 err = f.Truncate(0)
 121                 if err != nil {
 122                         log.Fatalf("truncate(%s): %s", cfg.PIDFile, err)
 123                 }
 124                 _, err = fmt.Fprint(f, os.Getpid())
 125                 if err != nil {
 126                         log.Fatalf("write(%s): %s", cfg.PIDFile, err)
 127                 }
 128                 err = f.Sync()
 129                 if err != nil {
 130                         log.Fatal("sync(%s): %s", cfg.PIDFile, err)
 131                 }
 132         }
 133
 134         if cfg.DefaultReplicas > 0 {
 135                 kc.Want_replicas = cfg.DefaultReplicas
 136         }
 137         kc.Client.Timeout = time.Duration(cfg.Timeout)
 138         go kc.RefreshServices(5*time.Minute, 3*time.Second)
 139
 140         listener, err = net.Listen("tcp", cfg.Listen)
 141         if err != nil {
 142                 log.Fatalf("listen(%s): %s", cfg.Listen, err)
 143         }
 144         if _, err := daemon.SdNotify(false, "READY=1"); err != nil {
 145                 log.Printf("Error notifying init daemon: %v", err)
 146         }
 147         log.Println("Listening at", listener.Addr())
 148
 149         // Shut down the server gracefully (by closing the listener)
 150         // if SIGTERM is received.
 151         term := make(chan os.Signal, 1)
 152         go func(sig <-chan os.Signal) {
 153                 s := <-sig
 154                 log.Println("caught signal:", s)
 155                 listener.Close()
 156         }(term)
 157         signal.Notify(term, syscall.SIGTERM)
 158         signal.Notify(term, syscall.SIGINT)
 159
 160         // Start serving requests.
 161         http.Serve(listener, MakeRESTRouter(!cfg.DisableGet, !cfg.DisablePut, kc))
 162
 163         log.Println("shutting down")
 164 }
 165
 166 type ApiTokenCache struct {
 167         tokens     map[string]int64
 168         lock       sync.Mutex
 169         expireTime int64
 170 }
 171
 172 // Cache the token and set an expire time.  If we already have an expire time
 173 // on the token, it is not updated.
 174 func (this *ApiTokenCache) RememberToken(token string) {
 175         this.lock.Lock()
 176         defer this.lock.Unlock()
 177
 178         now := time.Now().Unix()
 179         if this.tokens[token] == 0 {
 180                 this.tokens[token] = now + this.expireTime
 181         }
 182 }
 183
 184 // Check if the cached token is known and still believed to be valid.
 185 func (this *ApiTokenCache) RecallToken(token string) bool {
 186         this.lock.Lock()
 187         defer this.lock.Unlock()
 188
 189         now := time.Now().Unix()
 190         if this.tokens[token] == 0 {
 191                 // Unknown token
 192                 return false
 193         } else if now < this.tokens[token] {
 194                 // Token is known and still valid
 195                 return true
 196         } else {
 197                 // Token is expired
 198                 this.tokens[token] = 0
 199                 return false
 200         }
 201 }
 202
 203 func GetRemoteAddress(req *http.Request) string {
 204         if xff := req.Header.Get("X-Forwarded-For"); xff != "" {
 205                 return xff + "," + req.RemoteAddr
 206         }
 207         return req.RemoteAddr
 208 }
 209
 210 func CheckAuthorizationHeader(kc *keepclient.KeepClient, cache *ApiTokenCache, req *http.Request) (pass bool, tok string) {
 211         var auth string
 212         if auth = req.Header.Get("Authorization"); auth == "" {
 213                 return false, ""
 214         }
 215
 216         _, err := fmt.Sscanf(auth, "OAuth2 %s", &tok)
 217         if err != nil {
 218                 // Scanning error
 219                 return false, ""
 220         }
 221
 222         if cache.RecallToken(tok) {
 223                 // Valid in the cache, short circuit
 224                 return true, tok
 225         }
 226
 227         arv := *kc.Arvados
 228         arv.ApiToken = tok
 229         if err := arv.Call("HEAD", "users", "", "current", nil, nil); err != nil {
 230                 log.Printf("%s: CheckAuthorizationHeader error: %v", GetRemoteAddress(req), err)
 231                 return false, ""
 232         }
 233
 234         // Success!  Update cache
 235         cache.RememberToken(tok)
 236
 237         return true, tok
 238 }
 239
 240 type GetBlockHandler struct {
 241         *keepclient.KeepClient
 242         *ApiTokenCache
 243 }
 244
 245 type PutBlockHandler struct {
 246         *keepclient.KeepClient
 247         *ApiTokenCache
 248 }
 249
 250 type IndexHandler struct {
 251         *keepclient.KeepClient
 252         *ApiTokenCache
 253 }
 254
 255 type InvalidPathHandler struct{}
 256
 257 type OptionsHandler struct{}
 258
 259 // MakeRESTRouter
 260 //     Returns a mux.Router that passes GET and PUT requests to the
 261 //     appropriate handlers.
 262 //
 263 func MakeRESTRouter(
 264         enable_get bool,
 265         enable_put bool,
 266         kc *keepclient.KeepClient) *mux.Router {
 267
 268         t := &ApiTokenCache{tokens: make(map[string]int64), expireTime: 300}
 269
 270         rest := mux.NewRouter()
 271
 272         if enable_get {
 273                 rest.Handle(`/{locator:[0-9a-f]{32}\+.*}`,
 274                         GetBlockHandler{kc, t}).Methods("GET", "HEAD")
 275                 rest.Handle(`/{locator:[0-9a-f]{32}}`, GetBlockHandler{kc, t}).Methods("GET", "HEAD")
 276
 277                 // List all blocks
 278                 rest.Handle(`/index`, IndexHandler{kc, t}).Methods("GET")
 279
 280                 // List blocks whose hash has the given prefix
 281                 rest.Handle(`/index/{prefix:[0-9a-f]{0,32}}`, IndexHandler{kc, t}).Methods("GET")
 282         }
 283
 284         if enable_put {
 285                 rest.Handle(`/{locator:[0-9a-f]{32}\+.*}`, PutBlockHandler{kc, t}).Methods("PUT")
 286                 rest.Handle(`/{locator:[0-9a-f]{32}}`, PutBlockHandler{kc, t}).Methods("PUT")
 287                 rest.Handle(`/`, PutBlockHandler{kc, t}).Methods("POST")
 288                 rest.Handle(`/{any}`, OptionsHandler{}).Methods("OPTIONS")
 289                 rest.Handle(`/`, OptionsHandler{}).Methods("OPTIONS")
 290         }
 291
 292         rest.NotFoundHandler = InvalidPathHandler{}
 293
 294         return rest
 295 }
 296
 297 func SetCorsHeaders(resp http.ResponseWriter) {
 298         resp.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, POST, PUT, OPTIONS")
 299         resp.Header().Set("Access-Control-Allow-Origin", "*")
 300         resp.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Length, Content-Type, X-Keep-Desired-Replicas")
 301         resp.Header().Set("Access-Control-Max-Age", "86486400")
 302 }
 303
 304 func (this InvalidPathHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 305         log.Printf("%s: %s %s unroutable", GetRemoteAddress(req), req.Method, req.URL.Path)
 306         http.Error(resp, "Bad request", http.StatusBadRequest)
 307 }
 308
 309 func (this OptionsHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 310         log.Printf("%s: %s %s", GetRemoteAddress(req), req.Method, req.URL.Path)
 311         SetCorsHeaders(resp)
 312 }
 313
 314 var BadAuthorizationHeader = errors.New("Missing or invalid Authorization header")
 315 var ContentLengthMismatch = errors.New("Actual length != expected content length")
 316 var MethodNotSupported = errors.New("Method not supported")
 317
 318 var removeHint, _ = regexp.Compile("\\+K@[a-z0-9]{5}(\\+|$)")
 319
 320 func (this GetBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 321         SetCorsHeaders(resp)
 322
 323         locator := mux.Vars(req)["locator"]
 324         var err error
 325         var status int
 326         var expectLength, responseLength int64
 327         var proxiedURI = "-"
 328
 329         defer func() {
 330                 log.Println(GetRemoteAddress(req), req.Method, req.URL.Path, status, expectLength, responseLength, proxiedURI, err)
 331                 if status != http.StatusOK {
 332                         http.Error(resp, err.Error(), status)
 333                 }
 334         }()
 335
 336         kc := *this.KeepClient
 337
 338         var pass bool
 339         var tok string
 340         if pass, tok = CheckAuthorizationHeader(&kc, this.ApiTokenCache, req); !pass {
 341                 status, err = http.StatusForbidden, BadAuthorizationHeader
 342                 return
 343         }
 344
 345         // Copy ArvadosClient struct and use the client's API token
 346         arvclient := *kc.Arvados
 347         arvclient.ApiToken = tok
 348         kc.Arvados = &arvclient
 349
 350         var reader io.ReadCloser
 351
 352         locator = removeHint.ReplaceAllString(locator, "$1")
 353
 354         switch req.Method {
 355         case "HEAD":
 356                 expectLength, proxiedURI, err = kc.Ask(locator)
 357         case "GET":
 358                 reader, expectLength, proxiedURI, err = kc.Get(locator)
 359                 if reader != nil {
 360                         defer reader.Close()
 361                 }
 362         default:
 363                 status, err = http.StatusNotImplemented, MethodNotSupported
 364                 return
 365         }
 366
 367         if expectLength == -1 {
 368                 log.Println("Warning:", GetRemoteAddress(req), req.Method, proxiedURI, "Content-Length not provided")
 369         }
 370
 371         switch respErr := err.(type) {
 372         case nil:
 373                 status = http.StatusOK
 374                 resp.Header().Set("Content-Length", fmt.Sprint(expectLength))
 375                 switch req.Method {
 376                 case "HEAD":
 377                         responseLength = 0
 378                 case "GET":
 379                         responseLength, err = io.Copy(resp, reader)
 380                         if err == nil && expectLength > -1 && responseLength != expectLength {
 381                                 err = ContentLengthMismatch
 382                         }
 383                 }
 384         case keepclient.Error:
 385                 if respErr == keepclient.BlockNotFound {
 386                         status = http.StatusNotFound
 387                 } else if respErr.Temporary() {
 388                         status = http.StatusBadGateway
 389                 } else {
 390                         status = 422
 391                 }
 392         default:
 393                 status = http.StatusInternalServerError
 394         }
 395 }
 396
 397 var LengthRequiredError = errors.New(http.StatusText(http.StatusLengthRequired))
 398 var LengthMismatchError = errors.New("Locator size hint does not match Content-Length header")
 399
 400 func (this PutBlockHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 401         SetCorsHeaders(resp)
 402
 403         kc := *this.KeepClient
 404         var err error
 405         var expectLength int64
 406         var status = http.StatusInternalServerError
 407         var wroteReplicas int
 408         var locatorOut string = "-"
 409
 410         defer func() {
 411                 log.Println(GetRemoteAddress(req), req.Method, req.URL.Path, status, expectLength, kc.Want_replicas, wroteReplicas, locatorOut, err)
 412                 if status != http.StatusOK {
 413                         http.Error(resp, err.Error(), status)
 414                 }
 415         }()
 416
 417         locatorIn := mux.Vars(req)["locator"]
 418
 419         _, err = fmt.Sscanf(req.Header.Get("Content-Length"), "%d", &expectLength)
 420         if err != nil || expectLength < 0 {
 421                 err = LengthRequiredError
 422                 status = http.StatusLengthRequired
 423                 return
 424         }
 425
 426         if locatorIn != "" {
 427                 var loc *keepclient.Locator
 428                 if loc, err = keepclient.MakeLocator(locatorIn); err != nil {
 429                         status = http.StatusBadRequest
 430                         return
 431                 } else if loc.Size > 0 && int64(loc.Size) != expectLength {
 432                         err = LengthMismatchError
 433                         status = http.StatusBadRequest
 434                         return
 435                 }
 436         }
 437
 438         var pass bool
 439         var tok string
 440         if pass, tok = CheckAuthorizationHeader(&kc, this.ApiTokenCache, req); !pass {
 441                 err = BadAuthorizationHeader
 442                 status = http.StatusForbidden
 443                 return
 444         }
 445
 446         // Copy ArvadosClient struct and use the client's API token
 447         arvclient := *kc.Arvados
 448         arvclient.ApiToken = tok
 449         kc.Arvados = &arvclient
 450
 451         // Check if the client specified the number of replicas
 452         if req.Header.Get("X-Keep-Desired-Replicas") != "" {
 453                 var r int
 454                 _, err := fmt.Sscanf(req.Header.Get(keepclient.X_Keep_Desired_Replicas), "%d", &r)
 455                 if err == nil {
 456                         kc.Want_replicas = r
 457                 }
 458         }
 459
 460         // Now try to put the block through
 461         if locatorIn == "" {
 462                 if bytes, err := ioutil.ReadAll(req.Body); err != nil {
 463                         err = errors.New(fmt.Sprintf("Error reading request body: %s", err))
 464                         status = http.StatusInternalServerError
 465                         return
 466                 } else {
 467                         locatorOut, wroteReplicas, err = kc.PutB(bytes)
 468                 }
 469         } else {
 470                 locatorOut, wroteReplicas, err = kc.PutHR(locatorIn, req.Body, expectLength)
 471         }
 472
 473         // Tell the client how many successful PUTs we accomplished
 474         resp.Header().Set(keepclient.X_Keep_Replicas_Stored, fmt.Sprintf("%d", wroteReplicas))
 475
 476         switch err {
 477         case nil:
 478                 status = http.StatusOK
 479                 _, err = io.WriteString(resp, locatorOut)
 480
 481         case keepclient.OversizeBlockError:
 482                 // Too much data
 483                 status = http.StatusRequestEntityTooLarge
 484
 485         case keepclient.InsufficientReplicasError:
 486                 if wroteReplicas > 0 {
 487                         // At least one write is considered success.  The
 488                         // client can decide if getting less than the number of
 489                         // replications it asked for is a fatal error.
 490                         status = http.StatusOK
 491                         _, err = io.WriteString(resp, locatorOut)
 492                 } else {
 493                         status = http.StatusServiceUnavailable
 494                 }
 495
 496         default:
 497                 status = http.StatusBadGateway
 498         }
 499 }
 500
 501 // ServeHTTP implementation for IndexHandler
 502 // Supports only GET requests for /index/{prefix:[0-9a-f]{0,32}}
 503 // For each keep server found in LocalRoots:
 504 //   Invokes GetIndex using keepclient
 505 //   Expects "complete" response (terminating with blank new line)
 506 //   Aborts on any errors
 507 // Concatenates responses from all those keep servers and returns
 508 func (handler IndexHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 509         SetCorsHeaders(resp)
 510
 511         prefix := mux.Vars(req)["prefix"]
 512         var err error
 513         var status int
 514
 515         defer func() {
 516                 if status != http.StatusOK {
 517                         http.Error(resp, err.Error(), status)
 518                 }
 519         }()
 520
 521         kc := *handler.KeepClient
 522
 523         ok, token := CheckAuthorizationHeader(&kc, handler.ApiTokenCache, req)
 524         if !ok {
 525                 status, err = http.StatusForbidden, BadAuthorizationHeader
 526                 return
 527         }
 528
 529         // Copy ArvadosClient struct and use the client's API token
 530         arvclient := *kc.Arvados
 531         arvclient.ApiToken = token
 532         kc.Arvados = &arvclient
 533
 534         // Only GET method is supported
 535         if req.Method != "GET" {
 536                 status, err = http.StatusNotImplemented, MethodNotSupported
 537                 return
 538         }
 539
 540         // Get index from all LocalRoots and write to resp
 541         var reader io.Reader
 542         for uuid := range kc.LocalRoots() {
 543                 reader, err = kc.GetIndex(uuid, prefix)
 544                 if err != nil {
 545                         status = http.StatusBadGateway
 546                         return
 547                 }
 548
 549                 _, err = io.Copy(resp, reader)
 550                 if err != nil {
 551                         status = http.StatusBadGateway
 552                         return
 553                 }
 554         }
 555
 556         // Got index from all the keep servers and wrote to resp
 557         status = http.StatusOK
 558         resp.Write([]byte("\n"))
 559 }