diff --git a/internal/examples/fileserver/main.go b/internal/examples/fileserver/main.go index 34d1170..c2b8486 100644 --- a/internal/examples/fileserver/main.go +++ b/internal/examples/fileserver/main.go @@ -17,8 +17,10 @@ package main import ( "bytes" "context" + "errors" "flag" "html/template" + "io" "io/fs" "log" "mime" @@ -31,30 +33,41 @@ import ( testv1 "connectrpc.com/vanguard/internal/gen/vanguard/test/v1" "connectrpc.com/vanguard/internal/gen/vanguard/test/v1/testv1connect" "google.golang.org/genproto/googleapis/api/httpbody" + "google.golang.org/protobuf/types/known/emptypb" ) func main() { + if err := run(); err != nil { + log.Fatal(err) + } +} + +func run() error { flagset := flag.NewFlagSet("fileserver", flag.ExitOnError) port := flagset.String("p", "8100", "port to serve on") directory := flagset.String("d", ".", "the directory of static file to host") if err := flagset.Parse(os.Args[1:]); err != nil { - log.Fatal(err) + return err } - // Create Connect handler. - serviceHandler := &ContentService{ - FS: os.DirFS(*directory), + root, err := os.OpenRoot(*directory) + if err != nil { + return err } + defer root.Close() + + // Create Connect handler. + serviceHandler := &ContentService{root: root} // And wrap it with Vanguard. service := vanguard.NewService(testv1connect.NewContentServiceHandler(serviceHandler)) handler, err := vanguard.NewTranscoder([]*vanguard.Service{service}) if err != nil { - log.Fatal(err) + return err } // Now handler also supports REST requests, translated to Connect // using the HTTP annotations on the ContentService definition. log.Printf("Serving %s on HTTP port: %s\n", *directory, *port) - log.Fatal(http.ListenAndServe(":"+*port, handler)) + return http.ListenAndServe(":"+*port, handler) } var indexHTMLTemplate = template.Must(template.New("http").Parse(` @@ -78,7 +91,8 @@ var indexHTMLTemplate = template.Must(template.New("http").Parse(` type ContentService struct { testv1connect.UnimplementedContentServiceHandler - fs.FS + + root *os.Root } func (c *ContentService) Index(_ context.Context, req *connect.Request[testv1.IndexRequest]) (*connect.Response[httpbody.HttpBody], error) { @@ -88,7 +102,7 @@ func (c *ContentService) Index(_ context.Context, req *connect.Request[testv1.In name = "." } - file, err := c.Open(name) + file, err := c.root.Open(name) if err != nil { return nil, err } @@ -101,7 +115,7 @@ func (c *ContentService) Index(_ context.Context, req *connect.Request[testv1.In var data []byte if !stat.IsDir() { contentType = mime.TypeByExtension(filepath.Ext(name)) - data, err = fs.ReadFile(c.FS, name) + data, err = fs.ReadFile(c.root.FS(), name) if err != nil { return nil, err } @@ -113,7 +127,7 @@ func (c *ContentService) Index(_ context.Context, req *connect.Request[testv1.In Title: name, Files: make(map[string]string), } - entries, err := fs.ReadDir(c.FS, name) + entries, err := fs.ReadDir(c.root.FS(), name) if err != nil { return nil, err } @@ -132,3 +146,78 @@ func (c *ContentService) Index(_ context.Context, req *connect.Request[testv1.In Data: data, }), nil } + +// Upload receives a client-streaming RPC and writes the file to disk. +// +// Example with curl: +// +// curl -X POST --data-binary "@hello.txt" \ +// -H "Content-Type: application/octet-stream" \ +// "http://localhost:8100/hello.txt:upload" +func (c *ContentService) Upload( + _ context.Context, + stream *connect.ClientStream[testv1.UploadRequest], +) (*connect.Response[emptypb.Empty], error) { + var file *os.File + for stream.Receive() { + msg := stream.Msg() + if file == nil { + var err error + file, err = c.root.Create(msg.GetFilename()) + if err != nil { + return nil, err + } + defer file.Close() + log.Printf("Upload: %q", msg.GetFilename()) + } + if _, err := file.Write(msg.GetFile().GetData()); err != nil { + return nil, err + } + } + if err := stream.Err(); err != nil { + return nil, err + } + if file == nil { + return nil, connect.NewError(connect.CodeInvalidArgument, errors.New("no upload message received")) + } + return connect.NewResponse(&emptypb.Empty{}), nil +} + +// Download streams a file from disk as a server-streaming RPC. +// +// Example with curl: +// +// curl "http://localhost:8100/hello.txt:download" -o hello.txt +func (c *ContentService) Download( + _ context.Context, + req *connect.Request[testv1.DownloadRequest], + stream *connect.ServerStream[testv1.DownloadResponse], +) error { + file, err := c.root.Open(req.Msg.GetFilename()) + if err != nil { + return err + } + defer file.Close() + log.Printf("Download: %q", req.Msg.GetFilename()) + + buf := make([]byte, 32*1024) + for { + n, readErr := file.Read(buf) + if n > 0 { + if err := stream.Send(&testv1.DownloadResponse{ + File: &httpbody.HttpBody{ + ContentType: "application/octet-stream", + Data: buf[:n], + }, + }); err != nil { + return err + } + } + if readErr == io.EOF { + return nil + } + if readErr != nil { + return readErr + } + } +}