echoserver

server.go

package server

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"time"
)

// Server is a echo server.
type Server struct {
	listener *net.TCPListener
	timeout  time.Duration
	quit     chan bool
	exited   chan bool
}

func logf(format string, v ...interface{}) {
	log.Printf(format, v...)
}

// NewServer retruns Server with given listener and response.
func NewServer(port int, timeout time.Duration) (*Server, error) {
	tcpAddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("0.0.0.0:%d", port))
	if err != nil {
		return nil, err
	}
	l, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return nil, err
	}
	s := &Server{
		listener: l,
		quit:     make(chan bool),
		exited:   make(chan bool),
		timeout:  timeout,
	}
	return s, nil
}

// Start starts echo server.
func (server *Server) Start() {
	logf("Server Listening on %s", server.listener.Addr())

	var handlers sync.WaitGroup
	for {
		select {
		case <-server.quit:
			logf("Shutting down...")
			handlers.Wait()
			close(server.exited)
			return
		default:
			conn, err := server.listener.Accept()
			if err != nil {
				if opErr, ok := err.(*net.OpError); ok && opErr.Temporary() {
					continue
				}
				// socket has been closed.
				if strings.Contains(err.Error(), "use of closed network connection") {
					continue
				}
			}
			handlers.Add(1)
			go func() {
				if err := server.handleConnection(conn); err != nil {
					logf("handle error: %v", err)
				}
				handlers.Done()
			}()
		}
	}
}

// Stop stops server.
func (server *Server) Stop() error {
	logf("Server is stopping...")
	if err := server.listener.Close(); err != nil {
		return err
	}
	close(server.quit)
	<-server.exited
	logf("Server stopped successfully.")
	return nil
}

func (server *Server) handleConnection(conn net.Conn) error {
	logf("Connection accepted: %v -> %v", conn.RemoteAddr(), conn.LocalAddr())
	defer conn.Close()
	conn.SetDeadline(time.Now().Add(server.timeout))

	buf := make([]byte, 4*1024)
	for {
		n, err := conn.Read(buf)
		if err != nil {
			if opErr, ok := err.(*net.OpError); ok {
				if opErr.Timeout() {
					return err
				}
				if opErr.Temporary() {
					continue
				}
				return err
			}
			logf("read error: %v", err)
			return err
		}
		logf("read: %s", buf[:n])
		response := fmt.Sprintf("> %s", buf[:n])
		if _, err := conn.Write([]byte(response)); err != nil {
			logf("write error: %v", err)
			return nil
		}
		logf("write: %s", buf[:n])
	}
}