From 0eba1b126450dc6c62ded3542c825f645789a226 Mon Sep 17 00:00:00 2001
From: Sebastian Zagrodzki <sebek@zagrodzki.net>
Date: Fri, 15 Jan 2021 17:25:29 +0100
Subject: [PATCH] Checks for closed/uninitialized context and devices. (#93)

---
 device.go   |  2 +-
 usb.go      | 54 ++++++++++++++++++++++++++++++++++++++++++++---------
 usb_test.go |  5 +++++
 3 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/device.go b/device.go
index 2882899..88e360b 100644
--- a/device.go
+++ b/device.go
@@ -209,7 +209,7 @@ func (d *Device) Close() error {
 	if d.claimed != nil {
 		return fmt.Errorf("can't release the device %s, it has an open config %d", d, d.claimed.Desc.Number)
 	}
-	d.ctx.libusb.close(d.handle)
+	d.ctx.closeDev(d)
 	d.handle = nil
 	return nil
 }
diff --git a/usb.go b/usb.go
index c1bf678..159698a 100644
--- a/usb.go
+++ b/usb.go
@@ -125,11 +125,20 @@ see the excellent "USB in a nutshell" guide: http://www.beyondlogic.org/usbnutsh
 */
 package gousb
 
+import (
+	"errors"
+	"fmt"
+	"sync"
+)
+
 // Context manages all resources related to USB device handling.
 type Context struct {
 	ctx    *libusbContext
 	done   chan struct{}
 	libusb libusbIntf
+
+	mu      sync.Mutex
+	devices map[*Device]bool
 }
 
 // Debug changes the debug level. Level 0 means no debug, higher levels
@@ -146,9 +155,10 @@ func newContextWithImpl(impl libusbIntf) *Context {
 		panic(err)
 	}
 	ctx := &Context{
-		ctx:    c,
-		done:   make(chan struct{}),
-		libusb: impl,
+		ctx:     c,
+		done:    make(chan struct{}),
+		libusb:  impl,
+		devices: make(map[*Device]bool),
 	}
 	go impl.handleEvents(ctx.ctx, ctx.done)
 	return ctx
@@ -165,6 +175,9 @@ func NewContext() *Context {
 // If there are any errors enumerating the devices,
 // the final one is returned along with any successfully opened devices.
 func (c *Context) OpenDevices(opener func(desc *DeviceDesc) bool) ([]*Device, error) {
+	if c.ctx == nil {
+		return nil, errors.New("OpenDevices called on a closed or uninitialized Context")
+	}
 	list, err := c.libusb.getDevices(c.ctx)
 	if err != nil {
 		return nil, err
@@ -187,7 +200,11 @@ func (c *Context) OpenDevices(opener func(desc *DeviceDesc) bool) ([]*Device, er
 				reterr = err
 				continue
 			}
-			ret = append(ret, &Device{handle: handle, ctx: c, Desc: desc})
+			o := &Device{handle: handle, ctx: c, Desc: desc}
+			ret = append(ret, o)
+			c.mu.Lock()
+			c.devices[o] = true
+			c.mu.Unlock()
 		} else {
 			c.libusb.dereference(dev)
 		}
@@ -219,13 +236,32 @@ func (c *Context) OpenDeviceWithVIDPID(vid, pid ID) (*Device, error) {
 	return devs[0], nil
 }
 
+func (c *Context) closeDev(d *Device) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.libusb.close(d.handle)
+	delete(c.devices, d)
+}
+
+func (c *Context) checkOpenDevs() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if l := len(c.devices); l > 0 {
+		return fmt.Errorf("Context.Close called while %d Devices are still open, Close may be called only after all previously opened devices were successfuly closed.", l)
+	}
+	return nil
+}
+
 // Close releases the Context and all associated resources.
 func (c *Context) Close() error {
-	var ret error
-	c.done <- struct{}{}
-	if c.ctx != nil {
-		ret = c.libusb.exit(c.ctx)
+	if c.ctx == nil {
+		return nil
 	}
+	if err := c.checkOpenDevs(); err != nil {
+		return err
+	}
+	c.done <- struct{}{}
+	err := c.libusb.exit(c.ctx)
 	c.ctx = nil
-	return ret
+	return err
 }
diff --git a/usb_test.go b/usb_test.go
index e0c2c4e..d98f8bc 100644
--- a/usb_test.go
+++ b/usb_test.go
@@ -41,6 +41,11 @@ func TestOPenDevices(t *testing.T) {
 		t.Fatalf("OpenDevices(): %s", err)
 	}
 
+	// attempt to Close() should fail because of open devices
+	if err := c.Close(); err == nil {
+		t.Fatal("Context.Close succeeded while some devices were still open")
+	}
+
 	if got, want := len(devs), len(fakeDevices); got != want {
 		t.Fatalf("len(devs) = %d, want %d (based on num fake devs)", got, want)
 	}