whois42d/server.go

package main

//////////////////////////////////////////////////////////////////////////

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"path"
	"strconv"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
	"whois42d/whois"
)

//////////////////////////////////////////////////////////////////////////

type Server struct {
	registry         whois.Registry
	LastConnection   time.Time
	SocketActivation bool
	stopListening    int32
	activeWorkers    sync.WaitGroup
}

func New(opts options) *Server {
	registry := whois.New(opts.Datapath, opts.Header)
	return &Server{registry, time.Now(), false, 0, sync.WaitGroup{}}
}

func (s *Server) Run(listener *net.TCPListener) {
	atomic.StoreInt32(&s.stopListening, 0)
	s.activeWorkers.Add(1)
	defer s.activeWorkers.Done()
	defer listener.Close()
	for atomic.LoadInt32(&s.stopListening) != 1 {
		if e := listener.SetDeadline(time.Now().Add(time.Second)); e != nil {
			fmt.Fprintf(os.Stderr, "Error setting deadline: %v\n", e)
			continue
		}
		conn, err := listener.AcceptTCP()
		if err != nil {
			if err, ok := err.(net.Error); ok && err.Timeout() {
				continue
			} else {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(1)
			}
		}

		s.activeWorkers.Add(1)
		s.LastConnection = time.Now()
		go s.handleConn(conn)
	}
}

func (s *Server) Shutdown() {
	atomic.StoreInt32(&s.stopListening, 1)
	s.activeWorkers.Wait()
}

func (s *Server) handleConn(conn *net.TCPConn) {
	defer func() {
		conn.Close()
		s.activeWorkers.Done()
	}()

	s.registry.HandleQuery(conn)
}

//////////////////////////////////////////////////////////////////////////

type options struct {
	Port     uint   `json:"port"`
	Address  string `json:"address"`
	Registry string `json:"registry"`
	Datapath string
	Header   string `json:"header"`
}

func parseFlags() options {

	var o options

	// default options
	flag.UintVar(&o.Port, "port", 43, "port to listen")
	flag.StringVar(&o.Address, "address", "*", "address to listen")
	flag.StringVar(&o.Registry, "registry", ".", "path to dn42 registry")
	flag.StringVar(&o.Header,
		"header",
		"This is the dn42 whois query service.",
		"announcement header")

	flag.Parse()
	if o.Address == "*" {
		o.Address = ""
	}

	return o
}

//////////////////////////////////////////////////////////////////////////

func checkDataPath(registry string) (string, error) {
	dataPath := path.Join(registry, "data")

	if _, err := os.Stat(dataPath); err != nil {
		return "", fmt.Errorf("Cannot access '%s', should be in the registry repository: %s\n",
			dataPath,
			err)
	}
	return dataPath, nil
}

func createServer(opts options) (*Server, error) {
	dataPath, err := checkDataPath(opts.Registry)
	if err != nil {
		return nil, err
	}
	opts.Datapath = dataPath
	server := New(opts)

	address := opts.Address + ":" + strconv.Itoa(int(opts.Port))
	listener, err := net.Listen("tcp", address)
	if err != nil {
		return nil, err
	}
	go server.Run(listener.(*net.TCPListener))
	return server, nil
}

//////////////////////////////////////////////////////////////////////////
// server starts here

func main() {
	opts := parseFlags()

	// create TCP server
	server, err := createServer(opts)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}

	// Signals
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)
	signal.Notify(signals, syscall.SIGTERM)
	signal.Notify(signals, syscall.SIGINT)

	// wait for signal
	<-signals

	fmt.Printf("Shutting socket(s) down (takes up to 1s)\n")
	server.Shutdown()
}

//////////////////////////////////////////////////////////////////////////
// end of file