implement createInviteCode & createInviteCodes (#4)

Co-authored-by: hailey <[e-mail]>

.env.example

···
-COCOON_DID=
-COCOON_HOSTNAME=
+COCOON_DID="did:web:cocoon.example.com"
+COCOON_HOSTNAME="cocoon.example.com"
+COCOON_ROTATION_KEY_PATH="./rotation.key"
+COCOON_JWK_PATH="./jwk.key"
+COCOON_CONTACT_EMAIL="[e-mail]"
+COCOON_RELAYS=https://bsky.network
+# Generate with `openssl rand -hex 16`
+COCOON_ADMIN_PASSWORD=

README.md

···
 - [x] com.atproto.repo.applyWrites
 - [x] com.atproto.repo.createRecord
 - [x] com.atproto.repo.putRecord
-- [ ] com.atproto.repo.deleteRecord
+- [x] com.atproto.repo.deleteRecord
 - [x] com.atproto.repo.describeRepo
 - [x] com.atproto.repo.getRecord
 - [ ] com.atproto.repo.importRepo
···
 - [ ] com.atproto.server.checkAccountStatus
 - [x] com.atproto.server.confirmEmail
 - [x] com.atproto.server.createAccount
+- [x] com.atproto.server.createInviteCode
+- [x] com.atproto.server.createInviteCodes
 - [ ] com.atproto.server.deactivateAccount
 - [ ] com.atproto.server.deleteAccount
 - [x] com.atproto.server.deleteSession

cmd/cocoon/main.go

···
 				EnvVars:  []string{"COCOON_RELAYS"},
 			},
 			&cli.StringFlag{
+				Name:     "admin-password",
+				Required: true,
+				EnvVars:  []string{"COCOON_ADMIN_PASSWORD"},
+			},
+			&cli.StringFlag{
 				Name:     "smtp-user",
 				Required: false,
 				EnvVars:  []string{"COCOON_SMTP_USER"},
···
 		Version:   Version,
 	}
 
-	app.Run(os.Args)
+	if err := app.Run(os.Args); err != nil {
+		fmt.Printf("Error: %v\n", err)
+	}
 }
 
 var run = &cli.Command{
···
 			ContactEmail:    cmd.String("contact-email"),
 			Version:         Version,
 			Relays:          cmd.StringSlice("relays"),
+			AdminPassword:   cmd.String("admin-password"),
 			SmtpUser:        cmd.String("smtp-user"),
 			SmtpPass:        cmd.String("smtp-pass"),
 			SmtpHost:        cmd.String("smtp-host"),

server/handle_server_create_invite_code.go

···
 
 import (
 	"github.com/google/uuid"
+	"github.com/haileyok/cocoon/internal/helpers"
 	"github.com/haileyok/cocoon/models"
 	"github.com/labstack/echo/v4"
 )
 
+type ComAtprotoServerCreateInviteCodeRequest struct {
+	UseCount   int     `json:"useCount" validate:"required"`
+	ForAccount *string `json:"forAccount,omitempty"`
+}
+
+type ComAtprotoServerCreateInviteCodeResponse struct {
+	Code string `json:"code"`
+}
+
 func (s *Server) handleCreateInviteCode(e echo.Context) error {
-	ic := models.InviteCode{
-		Code: uuid.NewString(),
+	var req ComAtprotoServerCreateInviteCodeRequest
+	if err := e.Bind(&req); err != nil {
+		s.logger.Error("error binding", "error", err)
+		return helpers.ServerError(e, nil)
+	}
+
+	if err := e.Validate(req); err != nil {
+		s.logger.Error("error validating", "error", err)
+		return helpers.InputError(e, nil)
+	}
+
+	ic := uuid.NewString()
+
+	var acc string
+	if req.ForAccount == nil {
+		acc = "admin"
+	} else {
+		acc = *req.ForAccount
+	}
+
+	if err := s.db.Create(&models.InviteCode{
+		Code:              ic,
+		Did:               acc,
+		RemainingUseCount: req.UseCount,
+	}).Error; err != nil {
+		s.logger.Error("error creating invite code", "error", err)
+		return helpers.ServerError(e, nil)
 	}
 
-	return e.JSON(200, map[string]string{
-		"code": ic.Code,
+	return e.JSON(200, ComAtprotoServerCreateInviteCodeResponse{
+		Code: ic,
 	})
 }

server/handle_server_create_invite_codes.go

···
+package server
+
+import (
+	"github.com/Azure/go-autorest/autorest/to"
+	"github.com/google/uuid"
+	"github.com/haileyok/cocoon/internal/helpers"
+	"github.com/haileyok/cocoon/models"
+	"github.com/labstack/echo/v4"
+)
+
+type ComAtprotoServerCreateInviteCodesRequest struct {
+	CodeCount   *int      `json:"codeCount,omitempty"`
+	UseCount    int       `json:"useCount" validate:"required"`
+	ForAccounts *[]string `json:"forAccounts,omitempty"`
+}
+
+type ComAtprotoServerCreateInviteCodesResponse []ComAtprotoServerCreateInviteCodesItem
+
+type ComAtprotoServerCreateInviteCodesItem struct {
+	Account string   `json:"account"`
+	Codes   []string `json:"codes"`
+}
+
+func (s *Server) handleCreateInviteCodes(e echo.Context) error {
+	var req ComAtprotoServerCreateInviteCodesRequest
+	if err := e.Bind(&req); err != nil {
+		s.logger.Error("error binding", "error", err)
+		return helpers.ServerError(e, nil)
+	}
+
+	if err := e.Validate(req); err != nil {
+		s.logger.Error("error validating", "error", err)
+		return helpers.InputError(e, nil)
+	}
+
+	if req.CodeCount == nil {
+		req.CodeCount = to.IntPtr(1)
+	}
+
+	if req.ForAccounts == nil {
+		req.ForAccounts = to.StringSlicePtr([]string{"admin"})
+	}
+
+	var codes []ComAtprotoServerCreateInviteCodesItem
+
+	for _, did := range *req.ForAccounts {
+		var ics []string
+
+		for range *req.CodeCount {
+			ic := uuid.NewString()
+			ics = append(ics, ic)
+
+			if err := s.db.Create(&models.InviteCode{
+				Code:              ic,
+				Did:               did,
+				RemainingUseCount: req.UseCount,
+			}).Error; err != nil {
+				s.logger.Error("error creating invite code", "error", err)
+				return helpers.ServerError(e, nil)
+			}
+		}
+
+		codes = append(codes, ComAtprotoServerCreateInviteCodesItem{
+			Account: did,
+			Codes:   ics,
+		})
+	}
+
+	return e.JSON(200, codes)
+}

server/server.go

···
 	JwkPath         string
 	ContactEmail    string
 	Relays          []string
+	AdminPassword   string
 
 	SmtpUser  string
 	SmtpPass  string
···
 	ContactEmail   string
 	EnforcePeering bool
 	Relays         []string
+	AdminPassword  string
 	SmtpEmail      string
 	SmtpName       string
 }
···
 	return nil
 }
 
+func (s *Server) handleAdminMiddleware(next echo.HandlerFunc) echo.HandlerFunc {
+	return func(e echo.Context) error {
+		username, password, ok := e.Request().BasicAuth()
+		if !ok || username != "admin" || password != s.config.AdminPassword {
+			return helpers.InputError(e, to.StringPtr("Unauthorized"))
+		}
+
+		if err := next(e); err != nil {
+			e.Error(err)
+		}
+		
+		return nil
+	}
+}
+
 func (s *Server) handleSessionMiddleware(next echo.HandlerFunc) echo.HandlerFunc {
 	return func(e echo.Context) error {
 		authheader := e.Request().Header.Get("authorization")
···
 		return nil, fmt.Errorf("cocoon hostname must be set")
 	}
 
+	if args.AdminPassword == "" {
+		return nil, fmt.Errorf("admin password must be set")
+	}
+
 	if args.Logger == nil {
 		args.Logger = slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{}))
 	}
···
 			ContactEmail:   args.ContactEmail,
 			EnforcePeering: false,
 			Relays:         args.Relays,
+			AdminPassword:  args.AdminPassword,
 			SmtpName:       args.SmtpName,
 			SmtpEmail:      args.SmtpEmail,
 		},
···
 	// are there any routes that we should be allowing without auth? i dont think so but idk
 	s.echo.GET("/xrpc/*", s.handleProxy, s.handleSessionMiddleware)
 	s.echo.POST("/xrpc/*", s.handleProxy, s.handleSessionMiddleware)
+	
+	// admin routes
+	s.echo.POST("/xrpc/com.atproto.server.createInviteCode", s.handleCreateInviteCode, s.handleAdminMiddleware)
+	s.echo.POST("/xrpc/com.atproto.server.createInviteCodes", s.handleCreateInviteCodes, s.handleAdminMiddleware)
 }
 
 func (s *Server) Serve(ctx context.Context) error {