...

Source file src/edge-infra.dev/hack/tools/apk-repository/internal/server.go

Documentation: edge-infra.dev/hack/tools/apk-repository/internal

     1  package internal
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"net/http"
     8  	"strconv"
     9  	"strings"
    10  	"time"
    11  
    12  	"cloud.google.com/go/storage"
    13  	"github.com/go-logr/logr"
    14  )
    15  
    16  const (
    17  	ADDR   = "addr"
    18  	BYTES  = "bytes"
    19  	METHOD = "method"
    20  	RANGE  = "range"
    21  	STATUS = "status"
    22  	URL    = "url"
    23  )
    24  
    25  type BucketServer struct {
    26  	ctx          context.Context
    27  	router       *http.ServeMux
    28  	bucketHandle *storage.BucketHandle
    29  	logger       logr.Logger
    30  }
    31  
    32  func NewBucketServer(ctx context.Context, client *storage.Client, bucket string, logger logr.Logger) (*BucketServer, error) {
    33  	bucketHandle := client.Bucket(bucket)
    34  	_, err := bucketHandle.Attrs(ctx)
    35  	if err != nil {
    36  		return nil, fmt.Errorf("connection to the bucket %s could not be established", bucket)
    37  	}
    38  	logger.Info("bucket connection successful", "bucket", bucket)
    39  
    40  	bucketServer := &BucketServer{
    41  		ctx:          ctx,
    42  		bucketHandle: bucketHandle,
    43  		logger:       logger,
    44  	}
    45  
    46  	router := http.NewServeMux()
    47  	router.HandleFunc("/{object...}", bucketServer.getObject)
    48  	router.HandleFunc("/health", bucketServer.health)
    49  
    50  	return &BucketServer{
    51  		ctx:          ctx,
    52  		router:       router,
    53  		bucketHandle: bucketHandle,
    54  		logger:       logger,
    55  	}, nil
    56  }
    57  
    58  func (bs *BucketServer) Run(addr string) error {
    59  	server := &http.Server{
    60  		Addr:        addr,
    61  		Handler:     bs.router,
    62  		ReadTimeout: time.Second * 10,
    63  	}
    64  
    65  	bs.logger.Info("[apk-mirror] started", "bind", addr)
    66  
    67  	return server.ListenAndServe()
    68  }
    69  
    70  func (bs *BucketServer) handleError(w http.ResponseWriter, r *http.Request, err error) {
    71  	switch err {
    72  	case storage.ErrObjectNotExist:
    73  		http.Error(w, err.Error(), http.StatusNotFound)
    74  		bs.logger.Error(err, "object does not exist",
    75  			STATUS, strconv.Itoa(http.StatusNotFound),
    76  			ADDR, getEntityAddr(r),
    77  			METHOD, r.Method,
    78  			URL, r.URL.String(),
    79  		)
    80  	default:
    81  		http.Error(w, err.Error(), http.StatusInternalServerError)
    82  		bs.logger.Error(err, "error in proxy",
    83  			STATUS, strconv.Itoa(http.StatusInternalServerError),
    84  			ADDR, getEntityAddr(r),
    85  			METHOD, r.Method,
    86  			URL, r.URL.String(),
    87  		)
    88  	}
    89  }
    90  
    91  func (bs *BucketServer) getObjReader(w http.ResponseWriter, r *http.Request, obj *storage.ObjectHandle) (*storage.Reader, error) {
    92  	var objr *storage.Reader
    93  
    94  	ctx := r.Context()
    95  	attrs, err := obj.Attrs(ctx)
    96  	if err != nil {
    97  		return nil, err
    98  	}
    99  	setStrHeader(w, "Content-Type", attrs.ContentType)
   100  	setStrHeader(w, "Content-Language", attrs.ContentLanguage)
   101  	setStrHeader(w, "Cache-Control", attrs.CacheControl)
   102  	setStrHeader(w, "Content-Encoding", attrs.ContentEncoding)
   103  	setStrHeader(w, "Content-Disposition", attrs.ContentDisposition)
   104  
   105  	if r.Header.Get("Range") != "" { //nolint
   106  		byteRange := strings.Split(r.Header.Values("Range")[0], "=")[1]
   107  		rangeStart, err := strconv.Atoi(strings.Split(byteRange, "-")[0])
   108  		if err != nil {
   109  			return nil, err
   110  		}
   111  		rangeEnd, err := strconv.Atoi(strings.Split(byteRange, "-")[1])
   112  		if err != nil {
   113  			return nil, err
   114  		}
   115  		rangeLength := rangeEnd - rangeStart + 1
   116  		objr, err = obj.NewRangeReader(ctx, int64(rangeStart), int64(rangeLength))
   117  		if err != nil {
   118  			return nil, err
   119  		}
   120  	} else {
   121  		objr, err = obj.NewReader(ctx)
   122  		if err != nil {
   123  			return nil, err
   124  		}
   125  	}
   126  	return objr, nil
   127  }
   128  
   129  func (bs *BucketServer) health(w http.ResponseWriter, _ *http.Request) {
   130  	w.WriteHeader(http.StatusOK)
   131  }
   132  
   133  func (bs *BucketServer) getObject(w http.ResponseWriter, r *http.Request) {
   134  	objName := r.PathValue("object")
   135  	obj := bs.bucketHandle.Object(objName)
   136  	objr, err := bs.getObjReader(w, r, obj)
   137  	if err != nil {
   138  		bs.handleError(w, r, err)
   139  		return
   140  	}
   141  
   142  	bytesWritten, err := io.Copy(w, objr)
   143  	objr.Close()
   144  	setIntHeader(w, "Content-Length", bytesWritten)
   145  	if err != nil {
   146  		bs.handleError(w, r, err)
   147  		return
   148  	}
   149  
   150  	logPairs := []any{
   151  		STATUS, strconv.Itoa(http.StatusOK),
   152  		ADDR, getEntityAddr(r),
   153  		METHOD, r.Method,
   154  		URL, r.URL.String(),
   155  	}
   156  
   157  	if rangeVal := r.Header.Get("Range"); rangeVal != "" {
   158  		logPairs = append(
   159  			logPairs,
   160  			RANGE,
   161  			rangeVal,
   162  			BYTES,
   163  			strconv.FormatInt(bytesWritten, 10),
   164  		)
   165  	}
   166  
   167  	bs.logger.Info(fmt.Sprintf("got object %s", objName),
   168  		logPairs...,
   169  	)
   170  }
   171  
   172  func setStrHeader(w http.ResponseWriter, key string, value string) {
   173  	if value != "" {
   174  		w.Header().Add(key, value)
   175  	}
   176  }
   177  
   178  func setIntHeader(w http.ResponseWriter, key string, value int64) {
   179  	if value > 0 {
   180  		w.Header().Add(key, strconv.FormatInt(value, 10))
   181  	}
   182  }
   183  
   184  func getEntityAddr(r *http.Request) string {
   185  	if xForwardedFor := r.Header.Get("X-Forwarded-For"); xForwardedFor != "" {
   186  		return xForwardedFor
   187  	}
   188  
   189  	return r.RemoteAddr
   190  }
   191  

View as plain text