2 files changed, 78 insertions, 2 deletions

diff --git a/client/client.go b/client/client.go
index 773df54..7a6747b 100644
--- a/client/client.go
+++ b/client/client.go
@@ -49,6 +49,9 @@ type Client struct {
 	// cfg is the config
 	cfg *config.Config
 
+	// proxy is proxyguard
+	proxy Proxy
+
 	mu sync.Mutex
 }
 
@@ -486,6 +489,11 @@ func (c *Client) retrieveTokens(sid string, t srvtypes.Type) (*eduoauth.Token, e
 
 // Cleanup cleans up the VPN connection by sending a /disconnect
 func (c *Client) Cleanup(ck *cookie.Cookie) error {
+	// cleanup proxyguard
+	cerr := c.proxy.Cancel()
+	if cerr != nil {
+		log.Logger.Debugf("ProxyGuard cancel gave an error: %v", cerr)
+	}
 	srv, err := c.Servers.CurrentServer()
 	if err != nil {
 		return i18nerr.Wrap(err, "The current server was not found when cleaning up the connection")
diff --git a/client/proxy.go b/client/proxy.go
index e7757ca..849b524 100644
--- a/client/proxy.go
+++ b/client/proxy.go
@@ -1,7 +1,10 @@
 package client
 
 import (
+	"context"
 	"encoding/json"
+	"errors"
+	"sync"
 
 	"codeberg.org/eduVPN/proxyguard"
 
@@ -23,6 +26,70 @@ func (pl *ProxyLogger) Log(msg string) {
 	log.Logger.Infof("[Proxyguard] %s", msg)
 }
 
+// Proxy is a wrapper around ProxyGuard
+// that has the client
+// and a cancel for cancellation by common
+// and a mutex to protect against race conditions
+type Proxy struct {
+	c      *proxyguard.Client
+	mu     sync.Mutex
+	cancel context.CancelFunc
+}
+
+// NewClient creates a new ProxyGuard wrapper from client `c`
+func (p *Proxy) NewClient(c *proxyguard.Client) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.c = c
+}
+
+// Delete sets the inner client to nil
+func (p *Proxy) Delete() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.c = nil
+}
+
+// ErrNoProxyGuardCancel indicates that no ProxyGuard cancel function
+// was ever defined. You probably forgot to call `Tunnel`
+var ErrNoProxyGuardCancel = errors.New("no ProxyGuard cancel function")
+
+// Cancel cancels a running ProxyGuard tunnel
+// it returns an error if it cannot be canceled
+func (p *Proxy) Cancel() error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.cancel == nil {
+		return ErrNoProxyGuardCancel
+	}
+	p.cancel()
+	p.cancel = nil
+	return nil
+}
+
+// ErrNoProxyGuardClient is an error that is returned when no ProxyGuard client is created
+var ErrNoProxyGuardClient = errors.New("no ProxyGuard client created")
+
+// Tunnel is a wrapper around ProxyGuard tunnel that
+// that creates a new context that can be canceled
+func (p *Proxy) Tunnel(ctx context.Context, peer string) error {
+	p.mu.Lock()
+	if p.c == nil {
+		p.mu.Unlock()
+		return ErrNoProxyGuardClient
+	}
+	cctx, cf := context.WithCancel(ctx)
+	p.cancel = cf
+	p.mu.Unlock()
+	defer func() {
+		p.mu.Lock()
+		p.cancel = nil
+		p.mu.Unlock()
+	}()
+	// we set peer IPs to nil here as proxyguard already does a DNS request for us
+	return p.c.Tunnel(cctx, peer, nil)
+}
+
 // StartProxyguard starts proxyguard for proxied WireGuard connections
 func (c *Client) StartProxyguard(ck *cookie.Cookie, listen string, tcpsp int, peer string, gotFD func(fd int, pips string), ready func()) error {
 	var err error
@@ -45,8 +112,9 @@ func (c *Client) StartProxyguard(ck *cookie.Cookie, listen string, tcpsp int, pe
 		Ready: ready,
 	}
 
-	// we set peer IPs to nil here as proxyguard already does a DNS request for us
-	err = proxyc.Tunnel(ck.Context(), peer, nil)
+	c.proxy.NewClient(&proxyc)
+	defer c.proxy.Delete()
+	err = c.proxy.Tunnel(ck.Context(), peer)
 	if err != nil {
 		return i18nerr.Wrap(err, "The VPN proxy exited")
 	}