package suika import ( "context" "fmt" "io" "net" "sync" "time" "golang.org/x/time/rate" "gopkg.in/irc.v3" ) // ircConn is a generic IRC connection. It's similar to net.Conn but focuses on // reading and writing IRC messages. type ircConn interface { ReadMessage() (*irc.Message, error) WriteMessage(*irc.Message) error Close() error SetReadDeadline(time.Time) error SetWriteDeadline(time.Time) error RemoteAddr() net.Addr LocalAddr() net.Addr } func newNetIRCConn(c net.Conn) ircConn { type netConn net.Conn return struct { *irc.Conn netConn }{irc.NewConn(c), c} } type connOptions struct { Logger Logger RateLimitDelay time.Duration RateLimitBurst int } type conn struct { conn ircConn srv *Server logger Logger lock sync.Mutex outgoing chan<- *irc.Message closed bool closedCh chan struct{} } func newConn(srv *Server, ic ircConn, options *connOptions) *conn { outgoing := make(chan *irc.Message, 64) c := &conn{ conn: ic, srv: srv, outgoing: outgoing, logger: options.Logger, closedCh: make(chan struct{}), } go func() { ctx, cancel := c.NewContext(context.Background()) defer cancel() rl := rate.NewLimiter(rate.Every(options.RateLimitDelay), options.RateLimitBurst) for msg := range outgoing { if err := rl.Wait(ctx); err != nil { break } c.logger.Debugf("sent: %v", msg) c.conn.SetWriteDeadline(time.Now().Add(writeTimeout)) if err := c.conn.WriteMessage(msg); err != nil { c.logger.Printf("failed to write message: %v", err) break } } if err := c.conn.Close(); err != nil && !isErrClosed(err) { c.logger.Printf("failed to close connection: %v", err) } else { c.logger.Debugf("connection closed") } // Drain the outgoing channel to prevent SendMessage from blocking for range outgoing { // This space is intentionally left blank } }() c.logger.Debugf("new connection") return c } func (c *conn) isClosed() bool { c.lock.Lock() defer c.lock.Unlock() return c.closed } // Close closes the connection. It is safe to call from any goroutine. func (c *conn) Close() error { c.lock.Lock() defer c.lock.Unlock() if c.closed { return fmt.Errorf("connection already closed") } err := c.conn.Close() c.closed = true close(c.outgoing) close(c.closedCh) return err } func (c *conn) ReadMessage() (*irc.Message, error) { msg, err := c.conn.ReadMessage() if isErrClosed(err) { return nil, io.EOF } else if err != nil { return nil, err } c.logger.Debugf("received: %v", msg) return msg, nil } // SendMessage queues a new outgoing message. It is safe to call from any // goroutine. // // If the connection is closed before the message is sent, SendMessage silently // drops the message. func (c *conn) SendMessage(ctx context.Context, msg *irc.Message) { c.lock.Lock() defer c.lock.Unlock() if c.closed { return } select { case c.outgoing <- msg: // Success case <-ctx.Done(): c.logger.Printf("failed to send message: %v", ctx.Err()) } } func (c *conn) RemoteAddr() net.Addr { return c.conn.RemoteAddr() } func (c *conn) LocalAddr() net.Addr { return c.conn.LocalAddr() } // NewContext returns a copy of the parent context with a new Done channel. The // returned context's Done channel is closed when the connection is closed, // when the returned cancel function is called, or when the parent context's // Done channel is closed, whichever happens first. // // Canceling this context releases resources associated with it, so code should // call cancel as soon as the operations running in this Context complete. func (c *conn) NewContext(parent context.Context) (context.Context, context.CancelFunc) { ctx, cancel := context.WithCancel(parent) go func() { defer cancel() select { case <-ctx.Done(): // The parent context has been cancelled, or the caller has called // cancel() case <-c.closedCh: // The connection has been closed } }() return ctx, cancel }