优化部分代码

main.go

@@ -1,10 +1,13 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"log"
 	"net/http"
 	"os"
+	"os/signal"
+	"syscall"
 
 	"go.uber.org/zap"
 
@@ -27,21 +30,29 @@ func init() {
 
 func main() {
 	flag.Parse()
-	inject := debugInject()
+	call := logInject()
-	defer inject()
+	defer call()
 	loadConf()
 	gitea, err := providers.NewGitea(config.Auth.Server, config.Auth.Token)
 	if err != nil {
 		log.Fatalln(err)
 	}
-	server := pkg.NewPageServer(gitea, pkg.DefaultOptions(config.Domain))
+	giteaServer := pkg.NewPageServer(gitea, pkg.DefaultOptions(config.Domain))
-	mux := http.NewServeMux()
+	defer giteaServer.Close()
-	mux.Handle("/", server)
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
-	defer server.Close()
+	defer stop()
-	_ = http.ListenAndServe(config.Bind, mux)
+	svc := http.Server{Addr: config.Bind, Handler: giteaServer}
+	go func() {
+		select {
+		case <-ctx.Done():
+		}
+		zap.L().Debug("shutdown gracefully")
+		_ = svc.Close()
+	}()
+	_ = svc.ListenAndServe()
 }
 
-func debugInject() func() error {
+func logInject() func() error {
 	atom := zap.NewAtomicLevel()
 	if debug {
 		atom.SetLevel(zap.DebugLevel)

pkg/core/alias.go

@@ -0,0 +1,50 @@
+package core
+
+import (
+	"encoding/json"
+
+	"code.d7z.net/d7z-project/gitea-pages/pkg/utils"
+)
+
+type Alias struct {
+	Owner  string `json:"owner"`
+	Repo   string `json:"repo"`
+	Branch string `json:"branch"`
+}
+
+type DomainAlias struct {
+	config utils.Config
+}
+
+func NewDomainAlias(config utils.Config) *DomainAlias {
+	return &DomainAlias{config: config}
+}
+
+func (a *DomainAlias) Query(domain string) (*Alias, error) {
+	get, err := a.config.Get("alias/" + domain)
+	if err != nil {
+		return nil, err
+	}
+	rel := &Alias{}
+	if err = json.Unmarshal([]byte(get), rel); err != nil {
+		return nil, err
+	}
+	return rel, nil
+}
+
+func (a *DomainAlias) Bind(domain, owner, repo, branch string) error {
+	save := &Alias{
+		Owner:  owner,
+		Repo:   repo,
+		Branch: branch,
+	}
+	saveB, err := json.Marshal(save)
+	if err != nil {
+		return err
+	}
+	return a.config.Put("domain/"+domain, string(saveB), utils.TtlKeep)
+}
+
+func (a *DomainAlias) Unbind(domain string) error {
+	return a.config.Delete("domain/" + domain)
+}

pkg/core/backend_test.go

@@ -1,14 +0,0 @@
-package core
-
-import (
-	"encoding/json"
-	"testing"
-
-	"github.com/stretchr/testify/require"
-)
-
-func TestAny(t *testing.T) {
-	data := make(map[string]string)
-	err := json.Unmarshal([]byte(`{}`), &data)
-	require.NoError(t, err)
-}

pkg/core/domain.go

@@ -0,0 +1,98 @@
+package core
+
+import (
+	"fmt"
+	"os"
+	"regexp"
+	"strings"
+
+	"code.d7z.net/d7z-project/gitea-pages/pkg/utils"
+
+	"github.com/pkg/errors"
+	"go.uber.org/zap"
+)
+
+var portExp = regexp.MustCompile(`:\d+$`)
+
+type PageDomain struct {
+	*ServerMeta
+
+	alias         *DomainAlias
+	baseDomain    string
+	defaultBranch string
+}
+
+func NewPageDomain(meta *ServerMeta, config utils.Config, baseDomain, defaultBranch string) *PageDomain {
+	return &PageDomain{
+		baseDomain:    baseDomain,
+		defaultBranch: defaultBranch,
+		ServerMeta:    meta,
+		alias:         NewDomainAlias(config),
+	}
+}
+
+type PageDomainContent struct {
+	*PageMetaContent
+
+	Owner string
+	Repo  string
+	Path  string
+}
+
+func (m *PageDomainContent) CacheKey() string {
+	return fmt.Sprintf("%s/%s/%s/%s", m.Owner, m.Repo, m.CommitID, m.Path)
+}
+
+func (p *PageDomain) ParseDomainMeta(domain, path, branch string) (*PageDomainContent, error) {
+	if branch == "" {
+		branch = p.defaultBranch
+	}
+	domain = portExp.ReplaceAllString(strings.ToLower(domain), "")
+	pathS := strings.Split(strings.TrimPrefix(path, "/"), "/")
+
+	if !strings.HasSuffix(domain, "."+p.baseDomain) {
+		alias, err := p.alias.Query(domain)
+		if err != nil {
+			zap.L().Warn("未知域名", zap.String("base", p.baseDomain), zap.String("domain", domain))
+			return nil, os.ErrNotExist
+		}
+		zap.L().Debug("命中别名", zap.String("domain", domain), zap.Any("alias", alias))
+		return p.ReturnMeta(alias.Owner, alias.Repo, alias.Branch, pathS)
+	}
+	owner := strings.TrimSuffix(domain, "."+p.baseDomain)
+	repo := pathS[0]
+	if repo == "" {
+		// 回退到默认仓库
+		repo = p.baseDomain
+		zap.L().Debug("fail back to default repo", zap.String("repo", repo))
+	}
+	returnMeta, err := p.ReturnMeta(owner, repo, branch, pathS[1:])
+	if err != nil && !errors.Is(err, os.ErrNotExist) {
+		return nil, err
+	} else if err == nil {
+		return returnMeta, nil
+	}
+	// 回退到默认页面
+	return p.ReturnMeta(owner, repo, domain, pathS)
+}
+
+func (p *PageDomain) ReturnMeta(owner string, repo string, branch string, path []string) (*PageDomainContent, error) {
+	rel := &PageDomainContent{}
+	if meta, err := p.GetMeta(owner, repo, branch); err == nil {
+		rel.PageMetaContent = meta
+		rel.Owner = owner
+		rel.Repo = repo
+		rel.Path = strings.Join(path, "/")
+		if strings.HasSuffix(rel.Path, "/") || rel.Path == "" {
+			rel.Path = rel.Path + "index.html"
+		}
+		if meta.Domain != "" {
+			err = p.alias.Bind(meta.Domain, rel.Owner, rel.Repo, branch)
+			if err != nil {
+				zap.L().Warn("别名绑定失败", zap.Error(err))
+			}
+		}
+		return rel, nil
+	}
+	return nil, os.ErrNotExist
+}

pkg/core/meta.go

@@ -6,14 +6,19 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"regexp"
 	"strings"
 	"time"
 
+	"go.uber.org/zap"
+
 	"github.com/pkg/errors"
 
 	"code.d7z.net/d7z-project/gitea-pages/pkg/utils"
 )
 
+var regexpHostname = regexp.MustCompile(`^(?:([a-z0-9-]+|\*)\.)?([a-z0-9-]{1,61})\.([a-z0-9]{2,7})$`)
+
 type ServerMeta struct {
 	Backend
 
@@ -111,7 +116,12 @@ func (s *ServerMeta) GetMeta(owner, repo, branch string) (*PageMetaContent, erro
 	if cname, err := s.ReadString(owner, repo, rel.CommitID, "CNAME"); err != nil && !errors.Is(err, os.ErrNotExist) {
 		return nil, err
 	} else {
-		rel.Domain = strings.TrimSpace(cname)
+		cname = strings.TrimSpace(cname)
+		if regexpHostname.MatchString(cname) {
+			rel.Domain = cname
+		} else {
+			zap.L().Debug("指定的 CNAME 不合法", zap.String("cname", cname))
+		}
 	}
 	if find, _ := s.FileExists(owner, repo, rel.CommitID, ".history"); find {
 		rel.HistoryRouteMode = true

pkg/core/page.go

@@ -1,87 +0,0 @@
-package core
-
-import (
-	"fmt"
-	"os"
-	"regexp"
-	"strings"
-
-	"github.com/pkg/errors"
-	"go.uber.org/zap"
-)
-
-type PageDomain struct {
-	*ServerMeta
-
-	baseDomain    string
-	defaultBranch string
-}
-
-func NewPageDomain(meta *ServerMeta, baseDomain, defaultBranch string) *PageDomain {
-	return &PageDomain{
-		baseDomain:    baseDomain,
-		defaultBranch: defaultBranch,
-		ServerMeta:    meta,
-	}
-}
-
-type PageDomainContent struct {
-	*PageMetaContent
-
-	Owner string
-	Repo  string
-	Path  string
-}
-
-func (m *PageDomainContent) CacheKey() string {
-	return fmt.Sprintf("%s/%s/%s/%s", m.Owner, m.Repo, m.CommitID, m.Path)
-}
-
-func (p *PageDomain) ParseDomainMeta(domain, path, branch string) (*PageDomainContent, error) {
-	if branch == "" {
-		branch = p.defaultBranch
-	}
-	domain = regexp.MustCompile(`:\d+$`).ReplaceAllString(domain, "")
-
-	rel := &PageDomainContent{}
-	if !strings.HasSuffix(domain, "."+p.baseDomain) {
-		zap.L().Warn("Page Domain does not end with ."+p.baseDomain, zap.String("domain", domain))
-		return nil, os.ErrNotExist
-	}
-	rel.Owner = strings.TrimSuffix(domain, "."+p.baseDomain)
-	pathS := strings.Split(strings.TrimPrefix(path, "/"), "/")
-	rel.Repo = pathS[0]
-	defaultRepo := rel.Owner + "." + p.baseDomain
-	if rel.Repo == "" {
-		// 回退到默认仓库
-		rel.Repo = defaultRepo
-		zap.L().Debug("fail back to default repo", zap.String("repo", defaultRepo))
-	}
-
-	meta, err := p.GetMeta(rel.Owner, rel.Repo, branch)
-	if err != nil && !errors.Is(err, os.ErrNotExist) {
-		return nil, err
-	}
-	if err == nil {
-		rel.Path = strings.Join(pathS[1:], "/")
-		if strings.HasSuffix(rel.Path, "/") || rel.Path == "" {
-			rel.Path = rel.Path + "index.html"
-		}
-		rel.PageMetaContent = meta
-		return rel, nil
-	}
-	if defaultRepo == rel.Repo {
-		return nil, os.ErrNotExist
-	}
-	if meta, err := p.GetMeta(rel.Owner, defaultRepo, branch); err == nil {
-		rel.PageMetaContent = meta
-		rel.Repo = defaultRepo
-		rel.Path = strings.Join(pathS[1:], "/")
-		if strings.HasSuffix(rel.Path, "/") || rel.Path == "" {
-			rel.Path = rel.Path + "index.html"
-		}
-		return rel, nil
-	}
-
-	return nil, os.ErrNotExist
-}

pkg/server.go

@@ -8,12 +8,10 @@ import (
 	"path/filepath"
 	"time"
 
-	"github.com/pkg/errors"
-	"go.uber.org/zap"
-
 	"code.d7z.net/d7z-project/gitea-pages/pkg/core"
 	"code.d7z.net/d7z-project/gitea-pages/pkg/utils"
 	"github.com/pbnjay/memory"
+	"github.com/pkg/errors"
 )
 
 type ServerOptions struct {
@@ -49,7 +47,7 @@ type Server struct {
 func NewPageServer(backend core.Backend, options ServerOptions) *Server {
 	backend = core.NewCacheBackend(backend, options.Config, time.Minute)
 	svcMeta := core.NewServerMeta(options.HttpClient, backend, options.Config, time.Minute)
-	pageMeta := core.NewPageDomain(svcMeta, options.Domain, options.DefaultBranch)
+	pageMeta := core.NewPageDomain(svcMeta, options.Config, options.Domain, options.DefaultBranch)
 	reader := core.NewCacheBackendBlobReader(options.HttpClient, backend, options.Cache, options.MaxCacheSize)
 	return &Server{
 		meta:    pageMeta,
@@ -59,45 +57,36 @@ func NewPageServer(backend core.Backend, options ServerOptions) *Server {
 }
 
 func (s *Server) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
+	err := s.Serve(writer, request)
+	if errors.Is(err, os.ErrNotExist) {
+		http.Error(writer, "page not found.", http.StatusNotFound)
+	} else {
+		http.Error(writer, err.Error(), http.StatusInternalServerError)
+	}
+}
+
+func (s *Server) Serve(writer http.ResponseWriter, request *http.Request) error {
 	meta, err := s.meta.ParseDomainMeta(request.Host, request.URL.Path, request.URL.Query().Get("branch"))
 	if err != nil {
-		if errors.Is(err, os.ErrNotExist) {
+		return err
-			s.writeNotfoundError(writer, request.RequestURI)
-		} else {
-			s.writeError(writer, err)
-		}
-		return
 	}
 	result, err := s.reader.Open(meta.Owner, meta.Repo, meta.CommitID, meta.Path)
 	if err != nil {
-		// todo: 添加默认返回
+		return err
-		if errors.Is(err, os.ErrNotExist) {
-			s.writeNotfoundError(writer, request.RequestURI)
-		} else {
-			s.writeError(writer, err)
-		}
-		return
 	}
 	fileName := filepath.Base(meta.Path)
 	if reader, ok := result.(*utils.CacheContent); ok {
+		writer.Header().Add("X-Cache", "HIT")
 		http.ServeContent(writer, request, fileName, reader.LastModified, reader)
 		_ = reader.Close()
-		return
+	} else {
-	}
+		writer.Header().Add("X-Cache", "MISS")
 		writer.Header().Set("Content-Type", mime.TypeByExtension(meta.Path))
 		writer.WriteHeader(http.StatusOK)
 		_, _ = io.Copy(writer, result)
 		_ = result.Close()
-	return
 	}
-
+	return nil
-func (s *Server) writeError(writer http.ResponseWriter, err error) {
-	zap.L().Error("write error", zap.Error(err))
-	http.Error(writer, err.Error(), http.StatusInternalServerError)
-}
-
-func (s *Server) writeNotfoundError(writer http.ResponseWriter, path string) {
-	http.Error(writer, "page not found.", http.StatusNotFound)
 }
 
 func (s *Server) Close() error {

pkg/utils/cache.go

@@ -17,6 +17,14 @@ type CacheContent struct {
 	LastModified time.Time
 }
 
+func (c *CacheContent) ReadToString() (string, error) {
+	all, err := io.ReadAll(c)
+	if err != nil {
+		return "", err
+	}
+	return string(all), nil
+}
+
 type Cache interface {
 	Put(key string, reader io.Reader) error
 	// Get return CacheContent or nil when put nil io.reader

pkg/utils/config.go

@@ -6,8 +6,12 @@ import (
 	"os"
 	"sync"
 	"time"
+
+	"go.uber.org/zap"
 )
 
+const TtlKeep = -1
+
 type Config interface {
 	Put(key string, value string, ttl time.Duration) error
 	Get(key string) (string, error)
@@ -36,7 +40,7 @@ func NewConfigMemory(store string) (Config, error) {
 			}
 		}
 		for key, content := range item {
-			if time.Now().Before(content.ttl) {
+			if content.Ttl == nil || time.Now().Before(*content.Ttl) {
 				ret.data.Store(key, content)
 			}
 		}
@@ -46,14 +50,19 @@ func NewConfigMemory(store string) (Config, error) {
 }
 
 type configContent struct {
-	data string
+	Data string     `json:"data"`
-	ttl  time.Time
+	Ttl  *time.Time `json:"ttl,omitempty"`
 }
 
 func (m *ConfigMemory) Put(key string, value string, ttl time.Duration) error {
+	d := time.Now().Add(ttl)
+	td := &d
+	if ttl == -1 {
+		td = nil
+	}
 	m.data.Store(key, configContent{
-		data: value,
+		Data: value,
-		ttl:  time.Now().Add(ttl),
+		Ttl:  td,
 	})
 	return nil
 }
@@ -61,10 +70,10 @@ func (m *ConfigMemory) Put(key string, value string, ttl time.Duration) error {
 func (m *ConfigMemory) Get(key string) (string, error) {
 	if value, ok := m.data.Load(key); ok {
 		content := value.(configContent)
-		if time.Now().After(content.ttl) {
+		if content.Ttl != nil && time.Now().After(*content.Ttl) {
 			return "", os.ErrNotExist
 		}
-		return content.data, nil
+		return content.Data, nil
 	}
 	return "", os.ErrNotExist
 }
@@ -78,11 +87,16 @@ func (m *ConfigMemory) Close() error {
 	defer m.data.Clear()
 	if m.store != "" {
 		item := make(map[string]configContent)
+		now := time.Now()
 		m.data.Range(
 			func(key, value interface{}) bool {
-				item[key.(string)] = value.(configContent)
+				content := value.(configContent)
+				if content.Ttl == nil || now.Before(*content.Ttl) {
+					item[key.(string)] = content
+				}
 				return true
 			})
+		zap.L().Debug("回写内容到本地存储", zap.String("store", m.store), zap.Int("length", len(item)))
 		saved, err := json.Marshal(item)
 		if err != nil {
 			return err

pkg/utils/error.go

@@ -1,7 +0,0 @@
-package utils
-
-import "github.com/pkg/errors"
-
-type StackTracer interface {
-	StackTrace() errors.StackTrace
-}