Merge pull request #7 from opensciencegrid/addIpMapping

Add IP mapping

Author: djw8605

config/config.yaml

@@ -29,4 +29,15 @@ metrics:
 # the queue will be emptied.  The queue on disk is persistent between restarts.
 config_directory: /tmp/shoveler-queue
 
-
+# Mapping configuration
+# If map.all is set, all messages will be mapped to the configured origin.
+# For example, with the configuration
+# map:
+#   all: 172.0.0.4
+# If a packet comes in with the private ip address of 192.168.0.4, the packet origin will be changed to 172.0.0.4
+# The port is always preserved.
+# If you want multiple mappings, you can specify multiple map entries.
+# For example, with the configuration
+# map:
+#   192.168.0.5: 172.0.0.5
+#   192.168.0.6: 129.93.10.7

main.go

@@ -23,6 +23,10 @@ func main() {
 	} else {
 		log.SetLevel(log.WarnLevel)
 	}
+
+	// Configure the mapper
+	configureMap()
+
 	textFormatter := log.TextFormatter{}
 	textFormatter.DisableLevelTruncation = true
 	textFormatter.FullTimestamp = true
@@ -57,7 +61,12 @@ func main() {
 		panic(err)
 	}
 
-	defer conn.Close()
+	defer func(conn *net.UDPConn) {
+		err := conn.Close()
+		if err != nil {
+			log.Errorln("Error closing UDP connection:", err)
+		}
+	}(conn)
 
 	// Create the UDP forwarding destinations
 	var udpDestinations []net.Conn

map.go

@@ -0,0 +1,37 @@
+package main
+
+import (
+	"net"
+
+	"github.com/spf13/viper"
+)
+
+var (
+	mapAll string
+	ipMap  map[string]string
+)
+
+// configureMap sets the mapping configuration
+func configureMap() {
+	// First, check for the map environment variable
+	mapAll = viper.GetString("map.all")
+
+	// If the map is not set
+	ipMap = viper.GetStringMapString("map")
+
+}
+
+// mapIp returns the mapped IP address
+func mapIp(remote *net.UDPAddr) string {
+
+	if mapAll != "" {
+		return mapAll
+	}
+	if len(ipMap) == 0 {
+		return remote.IP.String()
+	}
+	if ip, ok := ipMap[remote.IP.String()]; ok {
+		return ip
+	}
+	return remote.IP.String()
+}

map_test.go

@@ -0,0 +1,52 @@
+package main
+
+import (
+	"bytes"
+	"github.com/spf13/viper"
+	"net"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSingleIp(t *testing.T) {
+	ip := net.UDPAddr{IP: net.ParseIP("192.168.0.5"), Port: 514}
+
+	// If the map is not set
+	config := Config{}
+	config.ReadConfig()
+	configureMap()
+	ipStr := mapIp(&ip)
+	assert.Equal(t, "192.168.0.5", ipStr, "Test when map is not set")
+
+	// If the map is set by environment variable
+	err := os.Setenv("SHOVELER_MAP_ALL", "172.168.0.5")
+	assert.NoError(t, err, "Failed to set environment variable SHOVELER_MAP_ALL")
+	config.ReadConfig()
+	configureMap()
+	ipStr = mapIp(&ip)
+	assert.Equal(t, "172.168.0.5", ipStr, "Test when map is set by environment variable")
+
+	// If the map is set by config file
+	err = os.Unsetenv("SHOVELER_MAP_ALL")
+	assert.NoError(t, err, "Failed to unset SHOVELER_MAP_ALL")
+	// any approach to require this configuration into your program.
+	var yamlExample = []byte(`
+map:
+  192.168.1.5: 172.168.1.6
+  172.168.2.7: 129.93.10.5
+`)
+	err = viper.ReadConfig(bytes.NewBuffer(yamlExample))
+	defer viper.Reset()
+	assert.NoError(t, err, "Failed to read config file")
+	configureMap()
+	defer func() {
+		ipMap = nil
+		mapAll = ""
+	}()
+	ip = net.UDPAddr{IP: net.ParseIP("192.168.1.5"), Port: 514}
+	assert.Equal(t, "172.168.1.6", mapIp(&ip), "Test when map is set by config file")
+	ip = net.UDPAddr{IP: net.ParseIP("172.168.2.7"), Port: 514}
+	assert.Equal(t, "129.93.10.5", mapIp(&ip), "Test when map is set by config file")
+}

packageudp.go

@@ -22,7 +22,7 @@ func packageUdp(packet []byte, remote *net.UDPAddr) []byte {
 	msg.Data = str
 
 	// add the remote
-	msg.Remote = remote.IP.String()
+	msg.Remote = mapIp(remote)
 	msg.Remote += ":" + strconv.Itoa(remote.Port)
 
 	msg.ShovelerVersion = version

packageudp_test.go

@@ -0,0 +1,56 @@
+package main
+
+import (
+	"encoding/json"
+	"github.com/stretchr/testify/assert"
+	"net"
+	"testing"
+)
+
+func TestPackageUdp(t *testing.T) {
+
+	// No mapping enabled
+	ip := net.UDPAddr{IP: net.ParseIP("192.168.0.7"), Port: 12345}
+	packaged := packageUdp([]byte("asdf"), &ip)
+	assert.NotEmpty(t, packaged)
+	// Parse back the json
+	var pkg Message
+	err := json.Unmarshal(packaged, &pkg)
+	assert.NoError(t, err)
+	assert.Equal(t, ip.String(), pkg.Remote, "Remote IP should be the same")
+	assert.Equal(t, "YXNkZg==", pkg.Data, "Data should be base64 encoded")
+}
+
+func TestPackageUdp_Mapping(t *testing.T) {
+	// Mapping enabled
+	ip := net.UDPAddr{IP: net.ParseIP("192.168.0.8"), Port: 12345}
+	mapAll = "172.0.0.9"
+	packaged := packageUdp([]byte("asdf"), &ip)
+	assert.NotEmpty(t, packaged)
+	// Parse back the json
+	var pkg Message
+	err := json.Unmarshal(packaged, &pkg)
+	assert.NoError(t, err)
+	assert.Equal(t, "172.0.0.9:12345", pkg.Remote, "Remote IP should be the same")
+	assert.Equal(t, "YXNkZg==", pkg.Data, "Data should be base64 encoded")
+	mapAll = ""
+}
+
+func TestPackageUdp_MappingMultiple(t *testing.T) {
+	// Mapping enabled
+	ip := net.UDPAddr{IP: net.ParseIP("192.168.0.8"), Port: 12345}
+	ipMap = make(map[string]string)
+	defer func() {
+		ipMap = nil
+	}()
+	ipMap["192.168.0.8"] = "172.0.0.10"
+	ipMap["192.168.0.9"] = "172.0.0.11"
+	packaged := packageUdp([]byte("asdf"), &ip)
+	assert.NotEmpty(t, packaged)
+	// Parse back the json
+	var pkg Message
+	err := json.Unmarshal(packaged, &pkg)
+	assert.NoError(t, err)
+	assert.Equal(t, "172.0.0.10:12345", pkg.Remote, "Remote IP should be the same")
+	assert.Equal(t, "YXNkZg==", pkg.Data, "Data should be base64 encoded")
+}