sshpoke/internal/server/manager.go

132 lines
3.3 KiB
Go

package server
import (
"context"
"errors"
"sync"
"github.com/Neur0toxine/sshpoke/internal/config"
"github.com/Neur0toxine/sshpoke/internal/docker"
"github.com/Neur0toxine/sshpoke/internal/logger"
"github.com/Neur0toxine/sshpoke/internal/server/driver"
"github.com/Neur0toxine/sshpoke/internal/server/driver/base"
"github.com/Neur0toxine/sshpoke/internal/server/driver/plugin"
"github.com/Neur0toxine/sshpoke/pkg/dto"
)
type Manager struct {
rw sync.RWMutex
servers map[string]base.Driver
plugins map[string]plugin.Plugin
statusMap map[string]serverStatus
statusLock sync.RWMutex
defaultServer string
}
type serverStatus struct {
Name string `json:"name"`
Connections Connections `json:"connections"`
}
var DefaultManager *Manager
var (
ErrNoServer = errors.New("server is not specified")
ErrNoSuchServer = errors.New("server does not exist")
)
func NewManager(ctx context.Context, servers []config.Server, defaultServer string) *Manager {
m := &Manager{
servers: make(map[string]base.Driver),
plugins: make(map[string]plugin.Plugin),
defaultServer: defaultServer,
}
for _, serverConfig := range servers {
server, err := driver.New(ctx, serverConfig.Name, serverConfig.Driver, serverConfig.Params)
if err != nil {
logger.Sugar.Errorf("cannot initialize server '%s': %s", serverConfig.Name, err)
continue
}
server.SetEventStatusCallback(m.eventStatusCallback(server.Name()))
if server.Driver() == config.DriverPlugin {
pl := server.(plugin.Plugin)
if pl.Token() == "" {
logger.Sugar.Warnf("server '%s' will not work because it doesn't have a token", pl.Name())
continue
}
existing, found := m.plugins[pl.Token()]
if found {
logger.Sugar.Fatalw("two plugins cannot have the same token",
"plugin1", existing.Name(), "plugin2", pl.Name(), "token", pl.Token())
}
m.plugins[pl.Token()] = pl
}
m.servers[serverConfig.Name] = server
}
return m
}
func (m *Manager) ProcessEvent(event dto.Event) error {
serverName := event.Container.Server
if serverName == "" {
serverName = m.defaultServer
}
if serverName == "" {
return ErrNoServer
}
defer m.rw.RUnlock()
m.rw.RLock()
srv, ok := m.servers[event.Container.Server]
if !ok {
return ErrNoSuchServer
}
return srv.Handle(event)
}
func (m *Manager) eventStatusCallback(serverName string) base.EventStatusCallback {
return func(status dto.EventStatus) {
m.processEventStatus(serverName, status)
}
}
func (m *Manager) processEventStatus(serverName string, event dto.EventStatus) {
m.statusLock.RLock()
_, exists := m.statusMap[serverName]
if !exists {
return
}
m.statusLock.RUnlock()
defer m.statusLock.Unlock()
m.statusLock.Lock()
item, found := docker.Default.GetContainer(event.ID)
if !found {
return
}
switch event.Type {
case dto.EventStart:
item.Domain = event.Domain
case dto.EventStop, dto.EventShutdown, dto.EventError:
item.Domain = ""
default:
return
}
m.statusMap[serverName].Connections[item.ID] = item
}
func (m *Manager) PluginByToken(token string) plugin.Plugin {
server, ok := m.plugins[token]
if !ok {
return nil
}
return server
}
func (m *Manager) WaitForShutdown() {
defer m.rw.RUnlock()
m.rw.RLock()
for _, srv := range m.servers {
srv.WaitForShutdown()
}
return
}