myip/server.go

package myip

import (
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/spf13/cobra"
)

var ServerCmd = &cobra.Command{
	Use:   "serve",
	Short: "Simple server that returns the resolved IP.",
	Run: func(cmd *cobra.Command, args []string) {
		if logToFile != "" {
			f, err := os.OpenFile(logToFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
			if err != nil {
				log.Fatalf("Error opening file: %v", err)
			}
			defer f.Close()
			mwriter := io.MultiWriter(f, os.Stderr)
			log.SetOutput(mwriter)
		}
		StartServer()
	},
}

var (
	// declared in client
	// port string

	addr string

	logToFile string
)

func init() {
	ServerCmd.Flags().StringVarP(&port, "port", "p", "5080", "server port to listen")
	ServerCmd.Flags().StringVarP(&addr, "address", "b", "0.0.0.0", "address the server binds on")

	ServerCmd.Flags().StringVarP(&logToFile, "log", "f", "", "log to file if filepath provided")
}

func StartServer() error {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		ip := GetIP(r)
		fmt.Fprint(w, ip+"\n")

		// logging
		user := r.URL.Query().Get("name")
		if user == "" {
			user = r.URL.String()
		}
		host, _, _ := net.SplitHostPort(r.RemoteAddr)
		log.Printf("%s %s %s %s %s\n", time.Now().Format("2006-01-02-15:04:05"), host, ip, user, r.Header.Get("User-Agent"))
	})

	err := http.ListenAndServe(addr+":"+port, simpleLogger(http.DefaultServeMux))
	if err != nil {
		log.Println(err)
		return err
	}
	return nil
}

func GetIP(r *http.Request) string {
	ip := r.Header.Get("X-Real-Ip")
	if ip == "" {
		ips := r.Header.Get("X-Forwarded-For")
		ipArr := strings.Split(ips, ",")
		ip = strings.Trim(ipArr[len(ipArr)-1], " ")
	}
	if ip == "" {
		var err error
		ip, _, err = net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			ip = r.RemoteAddr
		}
	}
	return ip
}

func simpleLogger(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

		// for name, values := range r.Header {
		// 	// Loop over all values for the name.
		// 	for _, value := range values {
		// 		fmt.Println(name, value)
		// 	}
		// }
		// log.Printf("%s %s %s %s\n", r.RemoteAddr, r.Method, r.URL, r.Header.Get("User-Agent"))
		handler.ServeHTTP(w, r)
	})
}