TCP Client and Server Implementation

package main

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// TCPServer represents a TCP server with graceful shutdown
type TCPServer struct {
	address string
	listener net.Listener
	wg      sync.WaitGroup
	ctx     context.Context
	cancel  context.CancelFunc
}

// NewTCPServer creates a new TCP server instance
func NewTCPServer(address string) *TCPServer {
	ctx, cancel := context.WithCancel(context.Background())
	return &TCPServer{
		address: address,
		ctx:     ctx,
		cancel:  cancel,
	}
}

// Start starts the TCP server and listens for connections
func (s *TCPServer) Start() error {
	// Create TCP listener
	listener, err := net.Listen("tcp", s.address)
	if err != nil {
		return fmt.Errorf("failed to start listener: %w", err)
	}
	s.listener = listener
	log.Printf("TCP server started on %s", s.address)

	// Handle OS signals for graceful shutdown
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigChan
		log.Println("Received shutdown signal, stopping server...")
		s.Stop()
	}()

	// Accept connections in a loop
	for {
		// Check if server is shutting down
		select {
		case <-s.ctx.Done():
			return nil
		default:
			// Accept new connection with timeout
			s.listener.(*net.TCPListener).SetDeadline(time.Now().Add(1 * time.Second))
			conn, err := s.listener.Accept()
			if err != nil {
				// Check if it's a timeout error (normal during shutdown)
				var netErr net.Error
				if errors.As(err, &netErr) && netErr.Timeout() {
					continue
				}
				// Check if listener is closed (normal during shutdown)
				if errors.Is(err, net.ErrClosed) {
					return nil
				}
				log.Printf("Failed to accept connection: %v", err)
				continue
			}

			// Handle connection in a separate goroutine
			s.wg.Add(1)
			go func() {
				defer s.wg.Done()
				s.handleConnection(conn)
			}()
		}
	}
}

// Stop gracefully stops the TCP server
func (s *TCPServer) Stop() {
	// Cancel context to signal shutdown
	s.cancel()

	// Close listener to stop accepting new connections
	if s.listener != nil {
		s.listener.Close()
	}

	// Wait for all connections to be handled
	s.wg.Wait()
	log.Println("TCP server stopped gracefully")
}

// handleConnection processes a single client connection
func (s *TCPServer) handleConnection(conn net.Conn) {
	// Ensure connection is closed when done
	defer func() {
		conn.Close()
		log.Printf("Connection closed: %s", conn.RemoteAddr())
	}()

	log.Printf("New connection from: %s", conn.RemoteAddr())

	// Set read/write deadlines to prevent hanging connections
	conn.SetReadDeadline(time.Now().Add(30 * time.Second))
	conn.SetWriteDeadline(time.Now().Add(10 * time.Second))

	// Create buffered reader for line-based communication
	reader := bufio.NewReader(conn)

	for {
		// Check if server is shutting down
		select {
		case <-s.ctx.Done():
			return
		default:
			// Read line from client
			message, err := reader.ReadString('\n')
			if err != nil {
				// Check for normal disconnection
				if errors.Is(err, net.ErrClosed) || err.Error() == "EOF" {
					return
				}
				log.Printf("Error reading from %s: %v", conn.RemoteAddr(), err)
				return
			}

			// Trim whitespace and process message
			message = message[:len(message)-1] // Remove newline
			log.Printf("Received from %s: %s", conn.RemoteAddr(), message)

			// Handle special commands
			switch message {
			case "quit":
				_, err := conn.Write([]byte("Goodbye!\n"))
				if err != nil {
					log.Printf("Error writing to %s: %v", conn.RemoteAddr(), err)
				}
				return
			case "ping":
				_, err := conn.Write([]byte("pong\n"))
				if err != nil {
					log.Printf("Error writing to %s: %v", conn.RemoteAddr(), err)
					return
				}
			default:
				// Echo message back to client
				response := fmt.Sprintf("Echo: %s\n", message)
				_, err := conn.Write([]byte(response))
				if err != nil {
					log.Printf("Error writing to %s: %v", conn.RemoteAddr(), err)
					return
				}
			}

			// Reset read deadline after successful read
			conn.SetReadDeadline(time.Now().Add(30 * time.Second))
		}
	}
}

// TCPClient represents a TCP client
type TCPClient struct {
	address string
	conn    net.Conn
}

// NewTCPClient creates a new TCP client instance
func NewTCPClient(address string) *TCPClient {
	return &TCPClient{address: address}
}

// Connect establishes a connection to the TCP server
func (c *TCPClient) Connect() error {
	conn, err := net.DialTimeout("tcp", c.address, 5*time.Second)
	if err != nil {
		return fmt.Errorf("failed to connect to server: %w", err)
	}
	c.conn = conn
	log.Printf("Connected to server at %s", c.address)
	return nil
}

// Send sends a message to the server and returns the response
func (c *TCPClient) Send(message string) (string, error) {
	if c.conn == nil {
		return "", errors.New("not connected to server")
	}

	// Set write deadline
	c.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))

	// Send message with newline (server expects line-based messages)
	_, err := c.conn.Write([]byte(message + "\n"))
	if err != nil {
		return "", fmt.Errorf("failed to send message: %w", err)
	}

	// Set read deadline
	c.conn.SetReadDeadline(time.Now().Add(5 * time.Second))

	// Read response
	reader := bufio.NewReader(c.conn)
	response, err := reader.ReadString('\n')
	if err != nil {
		return "", fmt.Errorf("failed to read response: %w", err)
	}

	// Remove newline and return
	return response[:len(response)-1], nil
}

// Close closes the client connection
func (c *TCPClient) Close() error {
	if c.conn != nil {
		return c.conn.Close()
	}
	return nil
}

func main() {
	// Start server in a goroutine
	server := NewTCPServer(":8080")
	go func() {
		if err := server.Start(); err != nil {
			log.Fatalf("Server error: %v", err)
		}
	}()

	// Give server time to start
	time.Sleep(100 * time.Millisecond)

	// Create and connect client
	client := NewTCPClient(":8080")
	if err := client.Connect(); err != nil {
		log.Fatalf("Client connection error: %v", err)
	}
	defer client.Close()

	// Send test messages
	messages := []string{"Hello Server!", "ping", "How are you?", "quit"}
	for _, msg := range messages {
		response, err := client.Send(msg)
		if err != nil {
			log.Printf("Error sending '%s': %v", msg, err)
			continue
		}
		fmt.Printf("Sent: %-15s Received: %s\n", msg, response)
	}

	// Give server time to process quit command
	time.Sleep(100 * time.Millisecond)

	// Stop server
	server.Stop()
}