1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
|
package http
import (
"context"
"fmt"
"net"
"net/http"
"strings"
"sync"
"time"
castorev1pb "code.tvl.fyi/tvix/castore-go"
storev1pb "code.tvl.fyi/tvix/store-go"
"github.com/go-chi/chi/middleware"
"github.com/go-chi/chi/v5"
log "github.com/sirupsen/logrus"
"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
)
type Server struct {
srv *http.Server
handler chi.Router
directoryServiceClient castorev1pb.DirectoryServiceClient
blobServiceClient castorev1pb.BlobServiceClient
pathInfoServiceClient storev1pb.PathInfoServiceClient
// When uploading NAR files to a HTTP binary cache, the .nar
// files are uploaded before the .narinfo files.
// We need *both* to be able to fully construct a PathInfo object.
// Keep a in-memory map of narhash(es) (in SRI) to (unnamed) root node and nar
// size.
// This is necessary until we can ask a PathInfoService for a node with a given
// narSha256.
narDbMu sync.Mutex
narDb map[string]*narData
}
type narData struct {
rootNode *castorev1pb.Node
narSize uint64
}
func New(
directoryServiceClient castorev1pb.DirectoryServiceClient,
blobServiceClient castorev1pb.BlobServiceClient,
pathInfoServiceClient storev1pb.PathInfoServiceClient,
enableAccessLog bool,
priority int,
) *Server {
r := chi.NewRouter()
r.Use(func(h http.Handler) http.Handler {
return otelhttp.NewHandler(h, "http.request")
})
if enableAccessLog {
r.Use(middleware.Logger)
}
r.Get("/", func(w http.ResponseWriter, r *http.Request) {
_, err := w.Write([]byte("nar-bridge"))
if err != nil {
log.Errorf("Unable to write response: %v", err)
}
})
r.Get("/nix-cache-info", func(w http.ResponseWriter, r *http.Request) {
_, err := w.Write([]byte(fmt.Sprintf("StoreDir: /nix/store\nWantMassQuery: 1\nPriority: %d\n", priority)))
if err != nil {
log.Errorf("Unable to write response: %v", err)
}
})
s := &Server{
handler: r,
directoryServiceClient: directoryServiceClient,
blobServiceClient: blobServiceClient,
pathInfoServiceClient: pathInfoServiceClient,
narDb: make(map[string]*narData),
}
registerNarPut(s)
registerNarinfoPut(s)
registerNarinfoGet(s)
registerNarGet(s)
return s
}
func (s *Server) Shutdown(ctx context.Context) error {
return s.srv.Shutdown(ctx)
}
// ListenAndServer starts the webserver, and waits for it being closed or
// shutdown, after which it'll return ErrServerClosed.
func (s *Server) ListenAndServe(addr string) error {
s.srv = &http.Server{
Handler: s.handler,
ReadTimeout: 500 * time.Second,
WriteTimeout: 500 * time.Second,
IdleTimeout: 500 * time.Second,
}
var listener net.Listener
var err error
// check addr. If it contains slashes, assume it's a unix domain socket.
if strings.Contains(addr, "/") {
listener, err = net.Listen("unix", addr)
} else {
listener, err = net.Listen("tcp", addr)
}
if err != nil {
return fmt.Errorf("unable to listen on %v: %w", addr, err)
}
return s.srv.Serve(listener)
}
|