implement single part range requests and HEAD requests

cmd/simple-S3-cache/main.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net/http"
+	"strconv"
 	"time"
 
 	"github.com/i5heu/simple-S3-cache/internal/config"
@@ -12,6 +13,7 @@ import (
 	"github.com/i5heu/simple-S3-cache/internal/ramCache"
 	"github.com/i5heu/simple-S3-cache/internal/storageCache"
 
+	"github.com/gotd/contrib/http_range"
 	"github.com/valyala/fasthttp"
 )
 
@@ -53,51 +55,131 @@ func (h *Handler) handler(ctx *fasthttp.RequestCtx) {
 	}()
 
 	ctx.Response.Header.Set("Access-Control-Allow-Origin", h.conf.CORSDomain)
-	ctx.Response.Header.Set("Access-Control-Allow-Methods", "GET")
+	ctx.Response.Header.Set("Access-Control-Allow-Methods", "GET, HEAD")
 	ctx.Response.Header.Set("Cache-Control", "max-age=31536000")
+	ctx.Response.Header.Set("Accept-Ranges", "bytes")
+
+	dataResult := h.getData(ctx)
+	if dataResult.Error != nil {
+		ctx.Response.SetStatusCode(500)
+		ctx.Response.SetBodyString(dataResult.Error.Error())
+		return
+	}
+	size = dataResult.Size
+	cached = dataResult.Cached
+
+	if ctx.Request.Header.Peek("Range") != nil {
+		// parse the range header
+		ranges, err := http_range.ParseRange(string(ctx.Request.Header.Peek("Range")), int64(size))
+		if err != nil {
+			ctx.Response.SetStatusCode(416)
+			return
+		}
+		// we only support one range
+		if len(ranges) > 1 {
+			ctx.Response.SetStatusCode(416)
+			return
+		}
+		// HTTP 206 (Partial Content)
+		ctx.Response.SetStatusCode(206)
+		ctx.Response.Header.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", ranges[0].Start, ranges[0].Start+ranges[0].Length, size))
+		ctx.Response.Header.Set("Content-Length", strconv.FormatUint(uint64(ranges[0].Length), 10))
+		ctx.Response.Header.Set("Content-Type", dataResult.MIME)
+	} else {
+		ctx.Response.Header.Set("Content-Length", strconv.FormatUint(uint64(size), 10))
+		ctx.Response.Header.Set("Content-Type", dataResult.MIME)
+	}
+
+	if ctx.IsGet() {
+		if ctx.Request.Header.Peek("Range") != nil {
+			// parse the range header
+			ranges, err := http_range.ParseRange(string(ctx.Request.Header.Peek("Range")), int64(size))
+			if err != nil {
+				ctx.Response.SetStatusCode(416)
+				return
+			}
+
+			ctx.Response.SetBody(dataResult.Data[ranges[0].Start : ranges[0].Start+ranges[0].Length])
+		} else {
+			ctx.Response.SetBody(dataResult.Data)
+		}
+
+	}
+}
+
+// size and error
+type DataStoreResult struct {
+	Data   []byte
+	MIME   string
+	Size   uint
+	Cached bool
+	Error  error
+}
+
+func (h *Handler) getData(ctx *fasthttp.RequestCtx) DataStoreResult {
 
 	url := config.GetCompleteURL(h.conf, string(ctx.Path()))
 	cachedData, mime := h.dataStoreRAM.GetCacheData(url)
 	if cachedData != nil {
-		cached = true
-		size = uint(len(cachedData))
-		ctx.Response.Header.Set("Content-Type", mime)
-		ctx.Response.SetBody(cachedData)
-		return
+		return DataStoreResult{
+			Data:   cachedData,
+			MIME:   mime,
+			Size:   uint(len(cachedData)),
+			Cached: true,
+			Error:  nil,
+		}
 	}
 
 	cachedData, mime = h.dataStoreStorage.GetCacheData(url)
 	if cachedData != nil {
-		cached = true
-		h.dataStoreRAM.CacheData(url, cachedData, mime)
-		size = uint(len(cachedData))
-		ctx.Response.Header.Set("Content-Type", mime)
-		ctx.Response.SetBody(cachedData)
-		return
+		return DataStoreResult{
+			Data:   cachedData,
+			MIME:   mime,
+			Size:   uint(len(cachedData)),
+			Cached: true,
+			Error:  nil,
+		}
 	}
 
 	res, err := http.Get(url)
+	defer res.Body.Close()
 	if err != nil {
 		fmt.Printf("error making http request: %s\n", err)
 		ctx.Response.SetStatusCode(500)
-		return
+		return DataStoreResult{
+			Data:   nil,
+			MIME:   "",
+			Size:   0,
+			Cached: false,
+			Error:  err,
+		}
 	}
-	defer res.Body.Close()
 
 	if res.StatusCode != 200 {
 		ctx.Response.SetStatusCode(res.StatusCode)
-		return
+		return DataStoreResult{
+			Data:   nil,
+			MIME:   "",
+			Size:   0,
+			Cached: false,
+			Error:  err,
+		}
 	}
 
 	bytes, err := ioutil.ReadAll(res.Body)
 	if err != nil {
 		fmt.Println("error reading response body: ", err)
 	}
 
-	size = uint(len(bytes))
 	sanitizedMime := helper.SanitizeMimeType(res.Header.Get("Content-Type"))
 	h.dataStoreRAM.CacheData(url, bytes, sanitizedMime)
 	h.dataStoreStorage.CacheData(url, bytes, sanitizedMime)
-	ctx.Response.Header.Set("Content-Type", sanitizedMime)
-	ctx.Response.SetBody(bytes)
+
+	return DataStoreResult{
+		Data:   bytes,
+		MIME:   sanitizedMime,
+		Size:   uint(len(bytes)),
+		Cached: false,
+		Error:  nil,
+	}
 }

go.mod

@@ -16,6 +16,7 @@ require (
 )
 
 require (
+	github.com/gotd/contrib v0.13.0
 	github.com/influxdata/influxdb-client-go/v2 v2.12.0
 	github.com/inhies/go-bytesize v0.0.0-20220417184213-4913239db9cf
 	github.com/mattn/go-colorable v0.1.12 // indirect

go.sum

@@ -18,6 +18,8 @@ github.com/go-openapi/swag v0.19.5/go.mod h1:POnQmlKehdgb5mhVOsnJFsivZCEZ/vjK9gh
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/golangci/lint-1 v0.0.0-20181222135242-d2cdd8c08219/go.mod h1:/X8TswGSh1pIozq4ZwCfxS0WA5JGXguxk94ar/4c87Y=
 github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
+github.com/gotd/contrib v0.13.0 h1:6Qq4jFAhq8bc8DyCM3sNiVj4TUClDLEEmUZR5Lr6L5M=
+github.com/gotd/contrib v0.13.0/go.mod h1:Y1vtBgGhfnLKXQTCE6HzU9+quvacOAC1zxXA1/I5t9w=
 github.com/influxdata/influxdb-client-go v1.4.0 h1:+KavOkwhLClHFfYcJMHHnTL5CZQhXJzOm5IKHI9BqJk=
 github.com/influxdata/influxdb-client-go v1.4.0/go.mod h1:S+oZsPivqbcP1S9ur+T+QqXvrYS3NCZeMQtBoH4D1dw=
 github.com/influxdata/influxdb-client-go/v2 v2.12.0 h1:LGct9uIp36IT+8RAJdmJGQbNonGi26YfYYSpDIyq8fI=