ea15af1833
- Fix YAML tags in auth config struct (json -> yaml) - Update CLI configs to use pre-hashed API keys - Remove double hashing in WebSocket client - Fix port mapping (9102 -> 9103) in CLI commands - Update permission keys to use jobs:read, jobs:create, etc. - Clean up all debug logging from CLI and server - All user roles now authenticate correctly: * Admin: Can queue jobs and see all jobs * Researcher: Can queue jobs and see own jobs * Analyst: Can see status (read-only access) Multi-user authentication is now fully functional.
89 lines
1.8 KiB
Go
Executable file
89 lines
1.8 KiB
Go
Executable file
// Package network provides SSH client and retry utilities.
|
|
package network
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/logging"
|
|
)
|
|
|
|
// SSHPool manages a pool of SSH client connections.
|
|
type SSHPool struct {
|
|
factory func() (*SSHClient, error)
|
|
pool chan *SSHClient
|
|
active int
|
|
maxConns int
|
|
mu sync.Mutex
|
|
logger *logging.Logger
|
|
}
|
|
|
|
// NewSSHPool creates a new SSH connection pool.
|
|
func NewSSHPool(maxConns int, factory func() (*SSHClient, error), logger *logging.Logger) *SSHPool {
|
|
return &SSHPool{
|
|
factory: factory,
|
|
pool: make(chan *SSHClient, maxConns),
|
|
maxConns: maxConns,
|
|
logger: logger,
|
|
}
|
|
}
|
|
|
|
// Get retrieves an SSH client from the pool or creates a new one.
|
|
func (p *SSHPool) Get(ctx context.Context) (*SSHClient, error) {
|
|
select {
|
|
case conn := <-p.pool:
|
|
return conn, nil
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
default:
|
|
p.mu.Lock()
|
|
if p.active < p.maxConns {
|
|
p.active++
|
|
p.mu.Unlock()
|
|
return p.factory()
|
|
}
|
|
p.mu.Unlock()
|
|
|
|
// Wait for available connection
|
|
select {
|
|
case conn := <-p.pool:
|
|
return conn, nil
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
}
|
|
}
|
|
}
|
|
|
|
// Put returns an SSH client to the pool.
|
|
func (p *SSHPool) Put(conn *SSHClient) {
|
|
select {
|
|
case p.pool <- conn:
|
|
default:
|
|
// Pool is full, close connection
|
|
err := conn.Close()
|
|
if err != nil {
|
|
p.logger.Warn("failed to close SSH connection", "error", err)
|
|
}
|
|
p.mu.Lock()
|
|
p.active--
|
|
p.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
// Close closes all connections in the pool.
|
|
func (p *SSHPool) Close() {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
|
|
// Close all connections in the pool
|
|
close(p.pool)
|
|
for conn := range p.pool {
|
|
err := conn.Close()
|
|
if err != nil {
|
|
p.logger.Warn("failed to close SSH connection", "error", err)
|
|
}
|
|
}
|
|
|
|
// Reset active count
|
|
p.active = 0
|
|
}
|