coze_studio/backend/main.go

/*
 * Copyright 2025 coze-dev Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Code generated by hertz generator.

package main

import (
	"context"
	"crypto/tls"
	"fmt"
	"os"
	"runtime/debug"
	"strings"

	"github.com/cloudwego/hertz/pkg/app/server"
	"github.com/cloudwego/hertz/pkg/common/config"
	"github.com/hertz-contrib/cors"
	"github.com/joho/godotenv"

	"github.com/coze-dev/coze-studio/backend/api/middleware"
	"github.com/coze-dev/coze-studio/backend/api/router"
	"github.com/coze-dev/coze-studio/backend/application"
	"github.com/coze-dev/coze-studio/backend/pkg/lang/conv"
	"github.com/coze-dev/coze-studio/backend/pkg/lang/ternary"
	"github.com/coze-dev/coze-studio/backend/pkg/logs"
	"github.com/coze-dev/coze-studio/backend/types/consts"
)

func main() {
	ctx := context.Background()
	// Please do not change the order of the function calls below
	setCrashOutput()

	if err := loadEnv(); err != nil {
		panic("loadEnv failed, err=" + err.Error())
	}

	setLogLevel()

	if err := application.Init(ctx); err != nil {
		panic("InitializeInfra failed, err=" + err.Error())
	}

	startHttpServer()
}

func startHttpServer() {
	maxRequestBodySize := os.Getenv("MAX_REQUEST_BODY_SIZE")
	maxSize := conv.StrToInt64D(maxRequestBodySize, 1024*1024*200)
	addr := getEnv("LISTEN_ADDR", ":8888")

	opts := []config.Option{
		server.WithHostPorts(addr),
		server.WithMaxRequestBodySize(int(maxSize)),
	}

	useSSL := getEnv(consts.UseSSL, "0")
	if useSSL == "1" {
		cert, err := tls.LoadX509KeyPair(getEnv(consts.SSLCertFile, ""),
			getEnv(consts.SSLKeyFile, ""))
		if err != nil {
			fmt.Println(err.Error())
		}
		cfg := &tls.Config{}
		cfg.Certificates = append(cfg.Certificates, cert)
		opts = append(opts, server.WithTLS(cfg))
		logs.Infof("Use SSL")
	}

	s := server.Default(opts...)

	// cors option
	config := cors.DefaultConfig()
	config.AllowAllOrigins = true
	config.AllowHeaders = []string{"*"}
	corsHandler := cors.New(config)

	// Middleware order matters
	s.Use(middleware.ContextCacheMW())     // must be first
	s.Use(middleware.RequestInspectorMW()) // must be second
	s.Use(middleware.SetHostMW())
	s.Use(middleware.SetLogIDMW())
	s.Use(corsHandler)
	s.Use(middleware.AccessLogMW())
	s.Use(middleware.OpenapiAuthMW())
	s.Use(middleware.SessionAuthMW())
	s.Use(middleware.I18nMW()) // must after SessionAuthMW

	router.GeneratedRegister(s)
	s.Spin()
}

func loadEnv() (err error) {
	appEnv := os.Getenv("APP_ENV")
	fileName := ternary.IFElse(appEnv == "", ".env", ".env."+appEnv)

	logs.Infof("load env file: %s", fileName)

	err = godotenv.Load(fileName)
	if err != nil {
		return fmt.Errorf("load env file(%s) failed, err=%w", fileName, err)
	}

	return err
}

func getEnv(key string, defaultValue string) string {
	v := os.Getenv(key)
	if v == "" {
		return defaultValue
	}
	return v
}

func setLogLevel() {
	level := strings.ToLower(os.Getenv("LOG_LEVEL"))

	logs.Infof("log level: %s", level)
	switch level {
	case "trace":
		logs.SetLevel(logs.LevelTrace)
	case "debug":
		logs.SetLevel(logs.LevelDebug)
	case "info":
		logs.SetLevel(logs.LevelInfo)
	case "notice":
		logs.SetLevel(logs.LevelNotice)
	case "warn":
		logs.SetLevel(logs.LevelWarn)
	case "error":
		logs.SetLevel(logs.LevelError)
	case "fatal":
		logs.SetLevel(logs.LevelFatal)
	default:
		logs.SetLevel(logs.LevelInfo)
	}
}

func setCrashOutput() {
	crashFile, _ := os.Create("crash.log")
	debug.SetCrashOutput(crashFile, debug.CrashOptions{})
}