File Upload Feature

internal/pkg/api/handleUpload.go

@@ -0,0 +1,261 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License;
+// you may not use this file except in compliance with the Elastic License.
+
+package api
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/elastic/fleet-server/v7/internal/pkg/bulk"
+	"github.com/elastic/fleet-server/v7/internal/pkg/cache"
+	"github.com/elastic/fleet-server/v7/internal/pkg/config"
+	"github.com/elastic/fleet-server/v7/internal/pkg/logger"
+	"github.com/elastic/fleet-server/v7/internal/pkg/uploader"
+	"github.com/elastic/fleet-server/v7/internal/pkg/uploader/cbor"
+	"github.com/elastic/go-elasticsearch/v8"
+	"github.com/julienschmidt/httprouter"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+)
+
+const (
+	// TODO: move to a config
+	maxFileSize    = 104857600 // 100 MiB
+	maxUploadTimer = 24 * time.Hour
+)
+
+func (rt Router) handleUploadStart(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+	start := time.Now()
+	reqID := r.Header.Get(logger.HeaderRequestID)
+	zlog := log.With().
+		Str(ECSHTTPRequestID, reqID).
+		Logger()
+
+	// authentication occurs inside here
+	// to check that key agent ID matches the ID in the body payload yet-to-be unmarshalled
+	if err := rt.ut.handleUploadStart(&zlog, w, r); err != nil {
+		writeUploadError(err, w, zlog, start, "error initiating upload process")
+		return
+	}
+}
+
+func (rt Router) handleUploadChunk(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+	start := time.Now()
+
+	id := ps.ByName("id")
+	chunkID := ps.ByName("num")
+
+	reqID := r.Header.Get(logger.HeaderRequestID)
+
+	zlog := log.With().
+		Str(LogAgentID, id).
+		Str(ECSHTTPRequestID, reqID).
+		Logger()
+
+	// simpler authentication check,  for high chunk throughput
+	// since chunk checksums must match transit hash
+	// AND optionally the initial hash, both having stricter auth checks
+	if _, err := authAPIKey(r, rt.bulker, rt.ut.cache); err != nil {
+		writeUploadError(err, w, zlog, start, "authentication failure for chunk write")
+		return
+	}
+
+	chunkNum, err := strconv.Atoi(chunkID)
+	if err != nil {
+		writeUploadError(err, w, zlog, start, "error parsing chunk index")
+		return
+	}
+	if err := rt.ut.handleUploadChunk(&zlog, w, r, id, chunkNum); err != nil {
+		writeUploadError(err, w, zlog, start, "error uploading chunk")
+		return
+	}
+}
+
+func (rt Router) handleUploadComplete(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+	start := time.Now()
+
+	id := ps.ByName("id")
+
+	reqID := r.Header.Get(logger.HeaderRequestID)
+
+	zlog := log.With().
+		Str(LogAgentID, id).
+		Str(ECSHTTPRequestID, reqID).
+		Logger()
+
+	// authentication occurs inside here, to ensure key agent ID
+	// matches the same agent ID the operation started with
+	if err := rt.ut.handleUploadComplete(&zlog, w, r, id); err != nil {
+		writeUploadError(err, w, zlog, start, "error finalizing upload")
+		return
+	}
+}
+
+type UploadT struct {
+	bulker      bulk.Bulk
+	chunkClient *elasticsearch.Client
+	cache       cache.Cache
+	uploader    *uploader.Uploader
+}
+
+func NewUploadT(cfg *config.Server, bulker bulk.Bulk, chunkClient *elasticsearch.Client, cache cache.Cache) *UploadT {
+	log.Info().
+		Interface("limits", cfg.Limits.ArtifactLimit).
+		Int64("maxFileSize", maxFileSize).
+		Msg("upload limits")
+
+	return &UploadT{
+		chunkClient: chunkClient,
+		bulker:      bulker,
+		cache:       cache,
+		uploader:    uploader.New(chunkClient, bulker, cache, maxFileSize, maxUploadTimer),
+	}
+}
+
+func (ut *UploadT) handleUploadStart(zlog *zerolog.Logger, w http.ResponseWriter, r *http.Request) error { //nolint:unparam // log is standard first arg for the handlers
+	// decode early to match agentID in the payload
+	payload, err := uploader.ReadDict(r.Body)
+	if err != nil {
+		if errors.Is(err, io.EOF) {
+			return fmt.Errorf("file info body is required: %w", err)
+		}
+		return err
+	}
+
+	// check API key matches payload agent ID
+	agentID, ok := payload.Str("agent_id")
+	if !ok || agentID == "" {
+		return errors.New("required field agent_id is missing")
+	}
+	if _, err := authAgent(r, &agentID, ut.bulker, ut.cache); err != nil {
+		return err
+	}
+
+	// validate payload, enrich with additional fields, and write metadata doc to ES
+	info, err := ut.uploader.Begin(r.Context(), payload)
+	if err != nil {
+		return err
+	}
+
+	// prepare and write response
+	out, err := json.Marshal(map[string]interface{}{
+		"upload_id":  info.ID,
+		"chunk_size": info.ChunkSize,
+	})
+	if err != nil {
+		return err
+	}
+	w.Header().Set("Content-Type", "application/json")
+	_, err = w.Write(out)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (ut *UploadT) handleUploadChunk(zlog *zerolog.Logger, w http.ResponseWriter, r *http.Request, uplID string, chunkID int) error {
+	chunkHash := strings.TrimSpace(r.Header.Get("X-Chunk-Sha2"))
+	if chunkHash == "" {
+		return errors.New("chunk hash header required")
+	}
+
+	upinfo, chunkInfo, err := ut.uploader.Chunk(r.Context(), uplID, chunkID, chunkHash)
+	if err != nil {
+		return err
+	}
+
+	// prevent over-sized chunks
+	data := http.MaxBytesReader(w, r.Body, uploader.MaxChunkSize)
+
+	// compute hash as we stream it
+	hash := sha256.New()
+	copier := io.TeeReader(data, hash)
+
+	ce := cbor.NewChunkWriter(copier, chunkInfo.Last, chunkInfo.BID, chunkInfo.SHA2, upinfo.ChunkSize)
+	if err := uploader.IndexChunk(r.Context(), ut.chunkClient, ce, upinfo.Source, chunkInfo.BID, chunkInfo.Pos); err != nil {
+		return err
+	}
+
+	hashsum := hex.EncodeToString(hash.Sum(nil))
+
+	if !strings.EqualFold(chunkHash, hashsum) {
+		// delete document, since we wrote it, but the hash was invalid
+		// context scoped to allow this operation to finish even if client disconnects
+		if err := uploader.DeleteChunk(context.Background(), ut.bulker, upinfo.Source, chunkInfo.BID, chunkInfo.Pos); err != nil {
+			zlog.Warn().Err(err).
+				Str("source", upinfo.Source).
+				Str("fileID", chunkInfo.BID).
+				Int("chunkNum", chunkInfo.Pos).
+				Msg("a chunk hash mismatch occurred, and fleet server was unable to remove the invalid chunk")
+		}
+		return uploader.ErrHashMismatch
+	}
+
+	return nil
+}
+
+func (ut *UploadT) handleUploadComplete(zlog *zerolog.Logger, w http.ResponseWriter, r *http.Request, uplID string) error { //nolint:unparam // log is standard first arg for the handlers
+	info, err := ut.uploader.GetUploadInfo(r.Context(), uplID)
+	if err != nil {
+		return err
+	}
+	// need to auth that it matches the ID in the initial
+	// doc, but that means we had to doc-lookup early
+	if _, err := authAgent(r, &info.AgentID, ut.bulker, ut.cache); err != nil {
+		return fmt.Errorf("Error authenticating for upload finalization: %w", err)
+	}
+
+	var req UploadCompleteRequest
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		return errors.New("unable to parse request body")
+	}
+
+	hash := strings.TrimSpace(req.TransitHash.SHA256)
+	if hash == "" {
+		return errors.New("transit hash required")
+	}
+
+	info, err = ut.uploader.Complete(r.Context(), uplID, hash)
+	if err != nil {
+		return err
+	}
+
+	_, err = w.Write([]byte(`{"status":"ok"}`))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// helper function for doing all the error responsibilities
+// at the HTTP edge
+func writeUploadError(err error, w http.ResponseWriter, zlog zerolog.Logger, start time.Time, msg string) {
+	cntUpload.IncError(err)
+	resp := NewHTTPErrResp(err)
+
+	zlog.WithLevel(resp.Level).
+		Err(err).
+		Int(ECSHTTPResponseCode, resp.StatusCode).
+		Int64(ECSEventDuration, time.Since(start).Nanoseconds()).
+		Msg(msg)
+	if e := resp.Write(w); e != nil {
+		zlog.Error().Err(e).Msg("failure writing error response")
+	}
+}
+
+type UploadCompleteRequest struct {
+	TransitHash struct {
+		SHA256 string `json:"sha256"`
+	} `json:"transithash"`
+}

internal/pkg/api/handleUpload_test.go

@@ -0,0 +1,8 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License;
+// you may not use this file except in compliance with the Elastic License.
+
+//go:build !integration
+// +build !integration
+
+package api

internal/pkg/api/metrics.go

@@ -33,6 +33,7 @@ var (
 	cntEnroll    routeStats
 	cntAcks      routeStats
 	cntStatus    routeStats
+	cntUpload    routeStats
 	cntArtifacts artifactStats
 )
 
@@ -104,6 +105,7 @@ func init() {
 	cntArtifacts.Register(routesRegistry.NewRegistry("artifacts"))
 	cntAcks.Register(routesRegistry.NewRegistry("acks"))
 	cntStatus.Register(routesRegistry.NewRegistry("status"))
+	cntUpload.Register(routesRegistry.NewRegistry("upload"))
 }
 
 func (rt *routeStats) IncError(err error) {

internal/pkg/api/router.go

@@ -25,11 +25,14 @@ import (
 )
 
 const (
-	RouteStatus    = "/api/status"
-	RouteEnroll    = "/api/fleet/agents/:id"
-	RouteCheckin   = "/api/fleet/agents/:id/checkin"
-	RouteAcks      = "/api/fleet/agents/:id/acks"
-	RouteArtifacts = "/api/fleet/artifacts/:id/:sha2"
+	RouteStatus         = "/api/status"
+	RouteEnroll         = "/api/fleet/agents/:id"
+	RouteCheckin        = "/api/fleet/agents/:id/checkin"
+	RouteAcks           = "/api/fleet/agents/:id/acks"
+	RouteArtifacts      = "/api/fleet/artifacts/:id/:sha2"
+	RouteUploadBegin    = "/api/fleet/uploads"
+	RouteUploadChunk    = "/api/fleet/uploads/:id/:num"
+	RouteUploadComplete = "/api/fleet/uploads/:id"
 )
 
 type Router struct {
@@ -41,12 +44,13 @@ type Router struct {
 	at     *ArtifactT
 	ack    *AckT
 	st     *StatusT
+	ut     *UploadT
 	sm     policy.SelfMonitor
 	tracer *apm.Tracer
 	bi     build.Info
 }
 
-func NewRouter(cfg *config.Server, bulker bulk.Bulk, ct *CheckinT, et *EnrollerT, at *ArtifactT, ack *AckT, st *StatusT, sm policy.SelfMonitor, tracer *apm.Tracer, bi build.Info) *Router {
+func NewRouter(cfg *config.Server, bulker bulk.Bulk, ct *CheckinT, et *EnrollerT, at *ArtifactT, ack *AckT, st *StatusT, ut *UploadT, sm policy.SelfMonitor, tracer *apm.Tracer, bi build.Info) *Router {
 	rt := &Router{
 		cfg:    cfg,
 		bulker: bulker,
@@ -57,6 +61,7 @@ func NewRouter(cfg *config.Server, bulker bulk.Bulk, ct *CheckinT, et *EnrollerT
 		ack:    ack,
 		st:     st,
 		tracer: tracer,
+		ut:     ut,
 		bi:     bi,
 	}
 
@@ -98,6 +103,21 @@ func (rt *Router) newHTTPRouter(addr string) *httprouter.Router {
 			RouteArtifacts,
 			limiter.WrapArtifact(rt.handleArtifacts, &cntArtifacts),
 		},
+		{
+			http.MethodPost,
+			RouteUploadBegin,
+			rt.handleUploadStart,
+		},
+		{
+			http.MethodPut,
+			RouteUploadChunk,
+			rt.handleUploadChunk,
+		},
+		{
+			http.MethodPost,
+			RouteUploadComplete,
+			rt.handleUploadComplete,
+		},
 	}
 
 	router := httprouter.New()

internal/pkg/api/router_test.go

@@ -48,7 +48,7 @@ func TestRun(t *testing.T) {
 	et, err := NewEnrollerT(verCon, cfg, nil, c)
 	require.NoError(t, err)
 
-	router := NewRouter(cfg, bulker, ct, et, nil, nil, nil, nil, nil, fbuild.Info{})
+	router := NewRouter(cfg, bulker, ct, et, nil, nil, nil, nil, nil, nil, fbuild.Info{})
 	errCh := make(chan error)
 
 	var wg sync.WaitGroup