Merge pull request #53 from giolekva/vpn

Basic in memory implementation of network Manager
diff --git a/core/vpn/engine/wireguard.go b/core/vpn/engine/wireguard.go
index 6059634..35e4747 100644
--- a/core/vpn/engine/wireguard.go
+++ b/core/vpn/engine/wireguard.go
@@ -1,7 +1,6 @@
 package engine
 
 import (
-	"encoding/hex"
 	"fmt"
 	"log"
 
@@ -160,9 +159,7 @@
 }
 
 func (e *WireguardEngine) DiscoEndpoint() string {
-	k := e.DiscoKey()
-	discoHex := hex.EncodeToString(k[:])
-	return fmt.Sprintf("%s%s", discoHex, controlclient.EndpointDiscoSuffix)
+	return e.DiscoKey().Endpoint()
 }
 
 func (e *WireguardEngine) Ping(ip netaddr.IP, cb func(*ipnstate.PingResult)) {
diff --git a/core/vpn/engine/wireguard_test.go b/core/vpn/engine/wireguard_test.go
index 39488e3..04927af 100644
--- a/core/vpn/engine/wireguard_test.go
+++ b/core/vpn/engine/wireguard_test.go
@@ -113,7 +113,8 @@
 		p := <-ping
 		if p.Err != "" {
 			t.Error(p.Err)
+		} else {
+			log.Printf("Ping received: %+v\n", p)
 		}
-		log.Printf("Ping received: %+v\n", p)
 	}
 }
diff --git a/core/vpn/in_memory_manager.go b/core/vpn/in_memory_manager.go
new file mode 100644
index 0000000..5efbfda
--- /dev/null
+++ b/core/vpn/in_memory_manager.go
@@ -0,0 +1,140 @@
+package vpn
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+
+	"github.com/giolekva/pcloud/core/vpn/types"
+)
+
+func errorDeviceNotFound(pubKey types.PublicKey) error {
+	return fmt.Errorf("Device not found: %s", pubKey)
+}
+
+type InMemoryManager struct {
+	lock         sync.Mutex
+	devices      []*types.DeviceInfo
+	keyToDevices map[types.PublicKey]*types.DeviceInfo
+	callbacks    map[types.PublicKey][]NetworkMapChangeCallback
+	ipm          IPManager
+}
+
+func NewInMemoryManager(ipm IPManager) Manager {
+	return &InMemoryManager{
+		devices:      make([]*types.DeviceInfo, 0),
+		keyToDevices: make(map[types.PublicKey]*types.DeviceInfo),
+		callbacks:    make(map[types.PublicKey][]NetworkMapChangeCallback),
+		ipm:          ipm,
+	}
+}
+
+func (m *InMemoryManager) RegisterDevice(d types.DeviceInfo) (*types.NetworkMap, error) {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	if _, ok := m.keyToDevices[d.PublicKey]; ok {
+		return nil, errors.New(fmt.Sprintf("Device with given public key is already registered: %s", d.PublicKey))
+	}
+	if _, err := m.ipm.New(d.PublicKey); err != nil {
+		return nil, err
+	}
+	m.keyToDevices[d.PublicKey] = &d
+	m.devices = append(m.devices, &d)
+	m.callbacks[d.PublicKey] = make([]NetworkMapChangeCallback, 0)
+	ret, err := m.genNetworkMap(&d)
+	if err != nil {
+		return nil, err
+	}
+	// TODO(giolekva): run this in a goroutine
+	for _, peer := range m.devices {
+		if peer.PublicKey != d.PublicKey {
+			netMap, err := m.genNetworkMap(peer)
+			if err != nil {
+				// TODO(giolekva): maybe return netmap of requested device anyways?
+				return nil, err
+			}
+			for _, cb := range m.callbacks[peer.PublicKey] {
+				cb(netMap)
+			}
+		}
+	}
+	return ret, nil
+}
+
+func (m *InMemoryManager) RemoveDevice(pubKey types.PublicKey) error {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	if _, ok := m.keyToDevices[pubKey]; !ok {
+		return errorDeviceNotFound(pubKey)
+	}
+	delete(m.keyToDevices, pubKey) // TODO(giolekva): maybe mark as deleted?
+	for i, peer := range m.devices {
+		if peer.PublicKey == pubKey {
+			m.devices[i] = m.devices[len(m.devices)-1]
+			m.devices = m.devices[:len(m.devices)-1]
+		}
+	}
+	for _, peer := range m.devices {
+		netMap, err := m.genNetworkMap(peer)
+		if err != nil {
+			return err
+		}
+		for _, cb := range m.callbacks[peer.PublicKey] {
+			cb(netMap)
+		}
+	}
+	return nil
+}
+
+func (m *InMemoryManager) GetNetworkMap(pubKey types.PublicKey) (*types.NetworkMap, error) {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	if d, ok := m.keyToDevices[pubKey]; ok {
+		return m.genNetworkMap(d)
+	}
+	return nil, errorDeviceNotFound(pubKey)
+}
+
+func (m *InMemoryManager) AddNetworkMapChangeCallback(pubKey types.PublicKey, cb NetworkMapChangeCallback) error {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	if _, ok := m.keyToDevices[pubKey]; ok {
+		m.callbacks[pubKey] = append(m.callbacks[pubKey], cb)
+	}
+	return errorDeviceNotFound(pubKey)
+}
+
+func (m *InMemoryManager) genNetworkMap(d *types.DeviceInfo) (*types.NetworkMap, error) {
+	vpnIP, err := m.ipm.Get(d.PublicKey)
+	// NOTE(giolekva): Should not happen as devices must have been already registered and assigned IP address.
+	// Maybe should return error anyways instead of panic?
+	if err != nil {
+		return nil, err
+	}
+	ret := types.NetworkMap{
+		Self: types.Node{
+			PublicKey:     d.PublicKey,
+			DiscoKey:      d.DiscoKey,
+			DiscoEndpoint: d.DiscoKey.Endpoint(),
+			IPPort:        d.IPPort,
+			VPNIP:         vpnIP,
+		},
+	}
+	for _, peer := range m.devices {
+		if d.PublicKey == peer.PublicKey {
+			continue
+		}
+		vpnIP, err := m.ipm.Get(peer.PublicKey)
+		if err != nil {
+			return nil, err
+		}
+		ret.Peers = append(ret.Peers, types.Node{
+			PublicKey:     peer.PublicKey,
+			DiscoKey:      peer.DiscoKey,
+			DiscoEndpoint: peer.DiscoKey.Endpoint(),
+			IPPort:        peer.IPPort,
+			VPNIP:         vpnIP,
+		})
+	}
+	return &ret, nil
+}
diff --git a/core/vpn/in_memory_manager_test.go b/core/vpn/in_memory_manager_test.go
new file mode 100644
index 0000000..880aacc
--- /dev/null
+++ b/core/vpn/in_memory_manager_test.go
@@ -0,0 +1,83 @@
+package vpn
+
+import (
+	"log"
+	"testing"
+
+	"inet.af/netaddr"
+	"tailscale.com/ipn/ipnstate"
+
+	"github.com/giolekva/pcloud/core/vpn/engine"
+	"github.com/giolekva/pcloud/core/vpn/types"
+)
+
+func TestTwoPeers(t *testing.T) {
+	ipm := NewSequentialIPManager(netaddr.MustParseIP("10.0.0.1"))
+	m := NewInMemoryManager(ipm)
+	privKeyA := types.NewPrivateKey()
+	a, err := engine.NewFakeWireguardEngine(12345, privKeyA)
+	if err != nil {
+		t.Fatal(err)
+	}
+	privKeyB := types.NewPrivateKey()
+	b, err := engine.NewFakeWireguardEngine(12346, privKeyB)
+	if err != nil {
+		t.Fatal(err)
+	}
+	nma, err := m.RegisterDevice(types.DeviceInfo{
+		privKeyA.Public(),
+		a.DiscoKey(),
+		netaddr.MustParseIPPort("127.0.0.1:12345"),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	m.AddNetworkMapChangeCallback(privKeyA.Public(), func(nm *types.NetworkMap) {
+		log.Printf("a: Received new NetworkMap: %+v\n", nm)
+		if err := a.Configure(nm); err != nil {
+			t.Fatal(err)
+		}
+	})
+	if err := a.Configure(nma); err != nil {
+		t.Fatal(err)
+	}
+	nmb, err := m.RegisterDevice(types.DeviceInfo{
+		privKeyB.Public(),
+		b.DiscoKey(),
+		netaddr.MustParseIPPort("127.0.0.1:12346"),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	m.AddNetworkMapChangeCallback(privKeyB.Public(), func(nm *types.NetworkMap) {
+		log.Printf("b: Received new NetworkMap: %+v\n", nm)
+		if err := b.Configure(nm); err != nil {
+			t.Fatal(err)
+		}
+	})
+	if err := b.Configure(nmb); err != nil {
+		t.Fatal(err)
+	}
+	ping := make(chan *ipnstate.PingResult, 2)
+	pingCb := func(p *ipnstate.PingResult) {
+		ping <- p
+	}
+	a.Ping(nmb.Self.VPNIP, pingCb)
+	b.Ping(nma.Self.VPNIP, pingCb)
+	for i := 0; i < 2; i++ {
+		p := <-ping
+		if p.Err != "" {
+			t.Error(p.Err)
+		} else {
+			log.Printf("Ping received: %+v\n", p)
+		}
+	}
+	if err := m.RemoveDevice(privKeyA.Public()); err != nil {
+		t.Fatal(err)
+	}
+	b.Ping(nma.Self.VPNIP, pingCb)
+	p := <-ping
+	if p.Err == "" {
+		t.Fatalf("Ping received even after removing device: %+v", p)
+	}
+}
diff --git a/core/vpn/ip_manager.go b/core/vpn/ip_manager.go
new file mode 100644
index 0000000..6758f3e
--- /dev/null
+++ b/core/vpn/ip_manager.go
@@ -0,0 +1,45 @@
+package vpn
+
+import (
+	"fmt"
+
+	"github.com/giolekva/pcloud/core/vpn/types"
+
+	"inet.af/netaddr"
+)
+
+// TODO(giolekva): Add Disable method which marks given IP as non-usable for future.
+// It will be used when devices get removed from the network, in which case IP should not be reused for safety reasons.
+type IPManager interface {
+	New(pubKey types.PublicKey) (netaddr.IP, error)
+	Get(pubKey types.PublicKey) (netaddr.IP, error)
+}
+
+type SequentialIPManager struct {
+	cur     netaddr.IP
+	keyToIP map[types.PublicKey]netaddr.IP
+}
+
+func NewSequentialIPManager(start netaddr.IP) IPManager {
+	return &SequentialIPManager{
+		cur:     start,
+		keyToIP: make(map[types.PublicKey]netaddr.IP),
+	}
+}
+
+func (m *SequentialIPManager) New(pubKey types.PublicKey) (netaddr.IP, error) {
+	ip := m.cur
+	if _, ok := m.keyToIP[pubKey]; ok {
+		return netaddr.IP{}, fmt.Errorf("Device with public key %s has already been assigned IP", pubKey)
+	}
+	m.keyToIP[pubKey] = ip
+	m.cur = m.cur.Next()
+	return ip, nil
+}
+
+func (m *SequentialIPManager) Get(pubKey types.PublicKey) (netaddr.IP, error) {
+	if ip, ok := m.keyToIP[pubKey]; ok {
+		return ip, nil
+	}
+	return netaddr.IP{}, fmt.Errorf("Device with public key %s pubKey does not have VPN IP assigned.", pubKey)
+}
diff --git a/core/vpn/ip_manager_test.go b/core/vpn/ip_manager_test.go
new file mode 100644
index 0000000..3a09569
--- /dev/null
+++ b/core/vpn/ip_manager_test.go
@@ -0,0 +1,52 @@
+package vpn
+
+import (
+	"log"
+	"testing"
+
+	"github.com/giolekva/pcloud/core/vpn/types"
+	"inet.af/netaddr"
+)
+
+func TestNewGet(t *testing.T) {
+	m := NewSequentialIPManager(netaddr.MustParseIP("10.0.0.1"))
+	a := types.NewPrivateKey()
+	b := types.NewPrivateKey()
+	ipA, err := m.New(a.Public())
+	if err != nil {
+		log.Fatal(err)
+	}
+	if ipA.String() != "10.0.0.1" {
+		t.Fatalf("Expected 10.0.0.1 Got: %s", ipA.String())
+	}
+	ipA, err = m.Get(a.Public())
+	if err != nil {
+		log.Fatal(err)
+	}
+	if ipA.String() != "10.0.0.1" {
+		t.Fatalf("Expected 10.0.0.1 Got: %s", ipA.String())
+	}
+	ipB, err := m.New(b.Public())
+	if err != nil {
+		log.Fatal(err)
+	}
+	if ipB.String() != "10.0.0.2" {
+		t.Fatalf("Expected 10.0.0.2 Got: %s", ipB.String())
+	}
+	ipB, err = m.Get(b.Public())
+	if err != nil {
+		log.Fatal(err)
+	}
+	if ipB.String() != "10.0.0.2" {
+		t.Fatalf("Expected 10.0.0.2 Got: %s", ipB.String())
+	}
+
+}
+
+func TestGetNonExistentPublicKey(t *testing.T) {
+	m := NewSequentialIPManager(netaddr.MustParseIP("10.0.0.1"))
+	a := types.NewPrivateKey()
+	if _, err := m.Get(a.Public()); err == nil {
+		t.Fatal("Returned IP for non existent public key")
+	}
+}
diff --git a/core/vpn/manager.go b/core/vpn/manager.go
new file mode 100644
index 0000000..78e8449
--- /dev/null
+++ b/core/vpn/manager.go
@@ -0,0 +1,27 @@
+package vpn
+
+import (
+	"github.com/giolekva/pcloud/core/vpn/types"
+)
+
+type NetworkMapChangeCallback func(*types.NetworkMap)
+
+// Manager interface manages mesh VPN configuration for all the devices registed by all users.
+// It does enforce device to device ACLs but delegates user authorization to the client.
+type Manager interface {
+	// Registers new device..
+	// Returns VPN network configuration on success and error otherwise.
+	// By default new devices have access to other machines owned by the same user
+	// and a PCloud entrypoint.
+	RegisterDevice(d types.DeviceInfo) (*types.NetworkMap, error)
+	// Completely removes device with given public key from the network.
+	RemoveDevice(pubKey types.PublicKey) error
+	// Returns network configuration for a device with a given public key.
+	// Result of this call must be encrypted with the same public key before
+	// sending it back to the client, so only the owner of it's corresponding
+	// private key is able to decrypt and use it.
+	GetNetworkMap(pubKey types.PublicKey) (*types.NetworkMap, error)
+	// AddNetworkMapChangeCallback can be used to receive new network configurations
+	// for a device with given public key.
+	AddNetworkMapChangeCallback(pubKey types.PublicKey, cb NetworkMapChangeCallback) error
+}
diff --git a/core/vpn/types/key.go b/core/vpn/types/key.go
index 1008634..5d02146 100644
--- a/core/vpn/types/key.go
+++ b/core/vpn/types/key.go
@@ -1,6 +1,12 @@
 package types
 
-import "tailscale.com/types/key"
+import (
+	"encoding/hex"
+	"fmt"
+
+	"tailscale.com/control/controlclient"
+	"tailscale.com/types/key"
+)
 
 // Generates new private key.
 func NewPrivateKey() PrivateKey {
@@ -11,3 +17,8 @@
 func (k PrivateKey) Public() PublicKey {
 	return PublicKey(key.Private(k).Public())
 }
+
+func (k DiscoKey) Endpoint() string {
+	discoHex := hex.EncodeToString(k[:])
+	return fmt.Sprintf("%s%s", discoHex, controlclient.EndpointDiscoSuffix)
+}
diff --git a/core/vpn/types/types.go b/core/vpn/types/types.go
index a8a1ed7..c421821 100644
--- a/core/vpn/types/types.go
+++ b/core/vpn/types/types.go
@@ -15,6 +15,12 @@
 //Public discovery key of the device.
 type DiscoKey wgcfg.Key
 
+type DeviceInfo struct {
+	PublicKey PublicKey
+	DiscoKey  DiscoKey
+	IPPort    netaddr.IPPort
+}
+
 // Represents single node in the network.
 type Node struct {
 	PublicKey     PublicKey