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") != "" {
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