ezshare/cmd/ezshare.go

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

}