package main import ( "context" "crypto/tls" "crypto/x509" "fmt" "io" "log" "net" "net/http" "os" "os/signal" "path/filepath" "time" "gitea.benny.dog/torjus/ezshare/certs" "gitea.benny.dog/torjus/ezshare/pb" "gitea.benny.dog/torjus/ezshare/server" "gitea.benny.dog/torjus/ezshare/store" "github.com/urfave/cli/v2" "google.golang.org/grpc" "google.golang.org/grpc/credentials" ) func main() { app := cli.App{ Name: "ezshare", Commands: []*cli.Command{ { Name: "serve", Usage: "Start ezshare server", Flags: []cli.Flag{ &cli.BoolFlag{ Name: "no-grpc", Usage: "Do not enable grpc.", }, &cli.BoolFlag{ Name: "no-http", Usage: "Do not enable http.", }, &cli.StringFlag{ Name: "grpc-addr", Usage: "Address to listen for grpc.", }, &cli.StringFlag{ Name: "http-addr", Usage: "Address to listen for http.", }, &cli.StringFlag{ Name: "hostname", Usage: "Hostname used in links", }, }, Action: ActionServe, }, { Name: "client", Usage: "Client commands", Flags: []cli.Flag{ &cli.StringFlag{ Name: "addr", Usage: "Address of server.", }, }, Subcommands: []*cli.Command{ { Name: "get", Usage: "Get file with id", ArgsUsage: "ID [ID]..", Action: ActionClientGet, }, { Name: "upload", Usage: "Upload file(s)", ArgsUsage: "PATH [PATH]..", Action: ActionClientUpload, }, }, }, { Name: "gencerts", Usage: "Generate certificates", Flags: []cli.Flag{ &cli.StringFlag{ Name: "out-dir", Usage: "Directory where certificates will be stored.", }, }, Action: ActionGencerts, }, }, } err := app.Run(os.Args) if err != nil { log.Printf("Error: %s\n", err) } } func ActionServe(c *cli.Context) error { fileStore := store.NewMemoryFileStore() // Setup shutdown-handling rootCtx, rootCancel := signal.NotifyContext(context.Background(), os.Interrupt) defer rootCancel() // Used to initiate grpc shutdown grpcCtx, grpcCancel := context.WithCancel(rootCtx) defer grpcCancel() // Cancelled once grpc is successfully shut down grpcShutdownCtx, grpcShutdownCancel := context.WithCancel(context.Background()) defer grpcShutdownCancel() // Start grpc server go func() { grpcAddr := ":50051" if c.IsSet("grpc-addr") { grpcAddr = c.String("grpc-addr") } grpcFileServer := server.NewGRPCFileServiceServer(fileStore) if c.IsSet("hostname") { grpcFileServer.Hostname = c.String("hostname") } lis, err := net.Listen("tcp", grpcAddr) if err != nil { log.Printf("Unable to setup grpc listener: %s\n", err) rootCancel() } srvCert, err := tls.X509KeyPair(certs.SrvCert, certs.SrvKey) if err != nil { log.Printf("%d %d", len(certs.SrvCert), len(certs.SrvKey)) log.Printf("Unable load server certs: %s\n", err) rootCancel() } certPool := x509.NewCertPool() if !certPool.AppendCertsFromPEM(certs.CACert) { log.Println("Unable to load CA cert") rootCancel() } tlsConfig := &tls.Config{ Certificates: []tls.Certificate{srvCert}, ClientAuth: tls.RequestClientCert, ClientCAs: certPool, } creds := credentials.NewTLS(tlsConfig) grpcServer := grpc.NewServer( grpc.Creds(creds), ) pb.RegisterFileServiceServer(grpcServer, grpcFileServer) // wait for cancel go func() { <-grpcCtx.Done() grpcServer.GracefulStop() }() log.Printf("Starting grpc server") if err = grpcServer.Serve(lis); err != nil { log.Printf("GRPC Shutdown with error: %s\n", err) rootCancel() } log.Println("GRPC Shutdown") grpcShutdownCancel() }() httpCtx, httpCancel := context.WithCancel(rootCtx) defer httpCancel() httpShutdownCtx, httpShutdownCancel := context.WithCancel(context.Background()) defer httpShutdownCancel() // Start http server go func() { httpAddr := ":8088" if c.IsSet("http-addr") { httpAddr = c.String("http-addr") } httpServer := server.NewHTTPSever(fileStore) httpServer.Addr = httpAddr // wait for cancel go func() { <-httpCtx.Done() timeoutCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() httpServer.Shutdown(timeoutCtx) }() log.Printf("Starting http server") if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed { log.Printf("HTTP Server shutdown with error: %s\n", err) rootCancel() } log.Println("HTTP Shutdown") httpShutdownCancel() }() <-grpcShutdownCtx.Done() <-httpShutdownCtx.Done() return nil } func ActionClientGet(c *cli.Context) error { addr := c.String("addr") conn, err := grpc.DialContext(c.Context, addr, grpc.WithInsecure()) if err != nil { return err } defer conn.Close() client := pb.NewFileServiceClient(conn) for _, arg := range c.Args().Slice() { req := &pb.GetFileRequest{Id: arg} resp, err := client.GetFile(c.Context, req) if err != nil { return err } filename := resp.File.FileId if resp.File.Metadata.OriginalFilename != "" { filename = resp.File.Metadata.OriginalFilename } f, err := os.Create(filename) if err != nil { return err } defer f.Close() if _, err := f.Write(resp.File.Data); err != nil { return err } fmt.Printf("Wrote file '%s'\n", filename) } return nil } func ActionClientUpload(c *cli.Context) error { addr := c.String("addr") clientCreds, err := getClientCreds() if err != nil { return err } conn, err := grpc.DialContext(c.Context, addr, grpc.WithTransportCredentials(clientCreds)) if err != nil { return err } defer conn.Close() client := pb.NewFileServiceClient(conn) for _, arg := range c.Args().Slice() { f, err := os.Open(arg) if err != nil { return err } data, err := io.ReadAll(f) if err != nil { return err } req := &pb.UploadFileRequest{Data: data, OriginalFilename: filepath.Base(arg)} resp, err := client.UploadFile(c.Context, req) if err != nil { return err } fmt.Printf("%s uploaded with id %s. Available at %s\n", arg, resp.Id, resp.FileUrl) } return nil } func ActionGencerts(c *cli.Context) error { return certs.GenCerts() } func getClientCreds() (credentials.TransportCredentials, error) { certPool := x509.NewCertPool() if !certPool.AppendCertsFromPEM(certs.CACert) { return nil, fmt.Errorf("unable to load ca cert") } clientCert, err := tls.X509KeyPair(certs.ClientCert, certs.ClientKey) if err != nil { return nil, fmt.Errorf("unable to load client cert: %s", err) } config := &tls.Config{ Certificates: []tls.Certificate{clientCert}, RootCAs: certPool, } return credentials.NewTLS(config), nil }