134 lines
3.4 KiB
Go
134 lines
3.4 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) {
|
|
logger.Sugar.Debugw("received EventStatus from server",
|
|
"serverName", serverName, "eventStatus", event)
|
|
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
|
|
}
|