// Command executorserver will starts a http server that receives command to run // programs inside a sandbox. package main import ( "context" "flag" "log" "net" "net/http" "os" "os/signal" "strconv" "strings" "time" "github.com/criyle/go-judge/env" "github.com/criyle/go-judge/filestore" "github.com/criyle/go-judge/pb" "github.com/criyle/go-judge/pkg/pool" "github.com/criyle/go-judge/worker" ginpprof "github.com/gin-contrib/pprof" "github.com/gin-gonic/gin" grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware" grpc_auth "github.com/grpc-ecosystem/go-grpc-middleware/auth" grpc_recovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery" grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus" ginprometheus "github.com/zsais/go-gin-prometheus" "golang.org/x/sync/errgroup" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" ) const ( envDebug = "DEBUG" envMetrics = "METRICS" envAddr = "HTTP_ADDR" envGRPC = "GRPC" envGRPCAddr = "GRPC_ADDR" envParallism = "PARALLISM" envToken = "TOKEN" ) var ( addr = flag.String("http", ":5050", "specifies the http binding address") grpcAddr = flag.String("grpc", ":5051", "specifies the grpc binding address") parallism = flag.Int("parallism", 4, "control the # of concurrency execution") tmpFsParam = flag.String("tmpfs", "size=16m,nr_inodes=4k", "tmpfs mount data (only for default mount with no mount.yaml)") dir = flag.String("dir", "", "specifies directory to store file upload / download (in memory by default)") silent = flag.Bool("silent", false, "do not print logs") netShare = flag.Bool("net", false, "do not unshare net namespace with host") mountConf = flag.String("mount", "mount.yaml", "specifics mount configuration file") cinitPath = flag.String("cinit", "", "container init absolute path") token = flag.String("token", "", "bearer token auth for REST / gRPC") printLog = func(v ...interface{}) {} work *worker.Worker ) func newFilsStore(dir string) filestore.FileStore { var fs filestore.FileStore if dir == "" { fs = filestore.NewFileMemoryStore() } else { os.MkdirAll(dir, 0755) fs = filestore.NewFileLocalStore(dir) } return fs } func initEnv() (bool, error) { eneableGRPC := false if s := os.Getenv(envAddr); s != "" { addr = &s } if os.Getenv(envGRPC) == "1" { eneableGRPC = true } if s := os.Getenv(envGRPCAddr); s != "" { eneableGRPC = true grpcAddr = &s } if s := os.Getenv(envParallism); s != "" { p, err := strconv.Atoi(s) if err != nil { return false, err } parallism = &p } if s := os.Getenv(envToken); s != "" { token = &s } return eneableGRPC, nil } func main() { flag.Parse() if !*silent { printLog = log.Println } enableGRPC, err := initEnv() if err != nil { log.Fatalln("init environment variable failed", err) } // Init environment pool fs := newFilsStore(*dir) b, err := env.NewBuilder(*cinitPath, *mountConf, *tmpFsParam, *netShare, printLog) if err != nil { log.Fatalln("create environment builder failed", err) } envPool := pool.NewPool(b) work = worker.New(fs, envPool, *parallism, *dir) work.Start() printLog("Starting worker with parallism", *parallism) var r *gin.Engine if *silent { gin.SetMode(gin.ReleaseMode) r = gin.New() r.Use(gin.Recovery()) } else { r = gin.Default() } // Metrics Handle if os.Getenv(envMetrics) == "1" { p := ginprometheus.NewPrometheus("gin") p.ReqCntURLLabelMappingFn = func(c *gin.Context) string { url := c.Request.URL.Path for _, p := range c.Params { if p.Key == "fid" { url = strings.Replace(url, p.Value, ":fid", 1) } } return url } p.Use(r) } // Add auth token if *token != "" { r.Use(tokenAuth(*token)) printLog("Attach token auth with token:", *token) } // File Handles fh := &fileHandle{fs: fs} r.GET("/file", fh.fileGet) r.POST("/file", fh.filePost) r.GET("/file/:fid", fh.fileIDGet) r.DELETE("/file/:fid", fh.fileIDDelete) // Run Handle r.POST("/run", handleRun) // WebSocket Handle r.GET("/ws", handleWS) // pprof if os.Getenv(envDebug) != "" { ginpprof.Register(r) } // gRPC server var grpcServer *grpc.Server if enableGRPC { streamMiddleware := []grpc.StreamServerInterceptor{ grpc_prometheus.StreamServerInterceptor, grpc_recovery.StreamServerInterceptor(), } unaryMiddleware := []grpc.UnaryServerInterceptor{ grpc_prometheus.UnaryServerInterceptor, grpc_recovery.UnaryServerInterceptor(), } if *token != "" { authFunc := grpcTokenAuth(*token) streamMiddleware = append(streamMiddleware, grpc_auth.StreamServerInterceptor(authFunc)) unaryMiddleware = append(unaryMiddleware, grpc_auth.UnaryServerInterceptor(authFunc)) } grpcServer = grpc.NewServer( grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(streamMiddleware...)), grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(unaryMiddleware...)), ) pb.RegisterExecutorServer(grpcServer, &execServer{fs: fs}) grpc_prometheus.Register(grpcServer) grpc_prometheus.EnableHandlingTimeHistogram() lis, err := net.Listen("tcp", *grpcAddr) if err != nil { log.Fatalln(err) } go func() { printLog("Starting grpc server at", *grpcAddr) printLog("GRPC serve", grpcServer.Serve(lis)) }() } srv := http.Server{ Addr: *addr, Handler: r, } go func() { printLog("Starting http server at", *addr) printLog("Http serve", srv.ListenAndServe()) }() // Graceful shutdown... sig := make(chan os.Signal, 1) signal.Notify(sig, os.Interrupt) <-sig printLog("Shutting Down...") ctx, cancel := context.WithTimeout(context.TODO(), time.Second*3) defer cancel() var eg errgroup.Group eg.Go(func() error { printLog("Http server shutdown") return srv.Shutdown(ctx) }) eg.Go(func() error { work.Shutdown() printLog("Worker shutdown") return nil }) if grpcServer != nil { eg.Go(func() error { grpcServer.GracefulStop() printLog("GRPC server shutdown") return nil }) } go func() { printLog("Shutdown Finished", eg.Wait()) cancel() }() <-ctx.Done() } func tokenAuth(token string) gin.HandlerFunc { const bearer = "Bearer " return func(c *gin.Context) { reqToken := c.GetHeader("Authorization") if strings.HasPrefix(reqToken, bearer) && reqToken[len(bearer):] == token { c.Next() return } c.AbortWithStatus(http.StatusUnauthorized) } } func grpcTokenAuth(token string) func(context.Context) (context.Context, error) { return func(ctx context.Context) (context.Context, error) { reqToken, err := grpc_auth.AuthFromMD(ctx, "bearer") if err != nil { return nil, err } if reqToken != token { return nil, status.Errorf(codes.Unauthenticated, "invalid auth token: %v", err) } return ctx, nil } }