Bugfix: allocate enough bytes for digital pins. Add more sanity check…

…s when allocating memory
DCC-EX · Jan 20, 2024 · fb414a7 · fb414a7
1 parent 818e05b
commit fb414a7
Showing 1 changed file with 54 additions and 34 deletions.
diff --git a/IO_EXIOExpander.h b/IO_EXIOExpander.h
@@ -1,5 +1,6 @@
 /*
  *  © 2022, Peter Cole. All rights reserved.
+ *  © 2024, Harald Barth. All rights reserved.
  *
  *  This file is part of EX-CommandStation
  *
@@ -98,34 +99,53 @@ class EXIOExpander : public IODevice {
           _numAnaloguePins = receiveBuffer[2];
 
           // See if we already have suitable buffers assigned
-          size_t digitalBytesNeeded = (_numDigitalPins + 7) / 8;
-          if (_digitalPinBytes < digitalBytesNeeded) {
-            // Not enough space, free any existing buffer and allocate a new one
-            if (_digitalPinBytes > 0) free(_digitalInputStates);
-            _digitalInputStates = (byte*) calloc(_digitalPinBytes, 1);
-            _digitalPinBytes = digitalBytesNeeded;
-          }
-          size_t analogueBytesNeeded = _numAnaloguePins * 2;
-          if (_analoguePinBytes < analogueBytesNeeded) {
-            // Free any existing buffers and allocate new ones.
-            if (_analoguePinBytes > 0) {
-              free(_analogueInputBuffer);
-              free(_analogueInputStates);
-              free(_analoguePinMap);
+          if (_numDigitalPins>0) {
+            size_t digitalBytesNeeded = (_numDigitalPins + 7) / 8;
+            if (_digitalPinBytes < digitalBytesNeeded) {
+              // Not enough space, free any existing buffer and allocate a new one
+              if (_digitalPinBytes > 0) free(_digitalInputStates);
+              if ((_digitalInputStates = (byte*) calloc(digitalBytesNeeded, 1)) != NULL) {
+		_digitalPinBytes = digitalBytesNeeded;
+	      } else {
+		DIAG(F("EX-IOExpander I2C:%s ERROR alloc %d bytes"), _I2CAddress.toString(), digitalBytesNeeded);
+		_deviceState = DEVSTATE_FAILED;
+		_digitalPinBytes = 0;
+		return;
+	      }
             }
-            _analogueInputStates = (uint8_t*) calloc(analogueBytesNeeded, 1);
-            _analogueInputBuffer = (uint8_t*) calloc(analogueBytesNeeded, 1);
-            _analoguePinMap = (uint8_t*) calloc(_numAnaloguePins, 1);
-            _analoguePinBytes = analogueBytesNeeded;
           }
-        } else {
-          DIAG(F("EX-IOExpander I2C:%s ERROR configuring device"), _I2CAddress.toString());
-          _deviceState = DEVSTATE_FAILED;
-          return;
-        }
-      } 
-      // We now need to retrieve the analogue pin map
-      if (status == I2C_STATUS_OK) {
+          if (_numAnaloguePins>0) {
+            size_t analogueBytesNeeded = _numAnaloguePins * 2;
+            if (_analoguePinBytes < analogueBytesNeeded) {
+              // Free any existing buffers and allocate new ones.
+              if (_analoguePinBytes > 0) {
+                free(_analogueInputBuffer);
+                free(_analogueInputStates);
+                free(_analoguePinMap);
+              }
+              _analogueInputStates = (uint8_t*) calloc(analogueBytesNeeded, 1);
+              _analogueInputBuffer = (uint8_t*) calloc(analogueBytesNeeded, 1);
+              _analoguePinMap = (uint8_t*) calloc(_numAnaloguePins, 1);
+	      if (_analogueInputStates  != NULL &&
+		  _analogueInputBuffer != NULL &&
+		  _analoguePinMap != NULL) {
+		_analoguePinBytes = analogueBytesNeeded;
+	      } else {
+		DIAG(F("EX-IOExpander I2C:%s ERROR alloc analog pin bytes"), _I2CAddress.toString());
+		_deviceState = DEVSTATE_FAILED;
+		_analoguePinBytes = 0;
+		return;
+	      }
+	    }
+	  }
+	} else {
+	  DIAG(F("EX-IOExpander I2C:%s ERROR configuring device"), _I2CAddress.toString());
+	  _deviceState = DEVSTATE_FAILED;
+	  return;
+	}
+      }
+      // We now need to retrieve the analogue pin map if there are analogue pins
+      if (status == I2C_STATUS_OK && _numAnaloguePins>0) {
         commandBuffer[0] = EXIOINITA;
         status = I2CManager.read(_I2CAddress, _analoguePinMap, _numAnaloguePins, commandBuffer, 1);
       }
@@ -239,15 +259,15 @@ class EXIOExpander : public IODevice {
 
     // If we're not doing anything now, check to see if a new input transfer is due.
     if (_readState == RDS_IDLE) {
-      if (currentMicros - _lastDigitalRead > _digitalRefresh) { // Delay for digital read refresh
+      if (_numDigitalPins>0 && currentMicros - _lastDigitalRead > _digitalRefresh) { // Delay for digital read refresh
         // Issue new read request for digital states.  As the request is non-blocking, the buffer has to
         // be allocated from heap (object state).
         _readCommandBuffer[0] = EXIORDD;
         I2CManager.read(_I2CAddress, _digitalInputStates, (_numDigitalPins+7)/8, _readCommandBuffer, 1, &_i2crb);
                                                                 // non-blocking read
         _lastDigitalRead = currentMicros;
         _readState = RDS_DIGITAL;
-      } else if (currentMicros - _lastAnalogueRead > _analogueRefresh) { // Delay for analogue read refresh
+      } else if (_numAnaloguePins>0 && currentMicros - _lastAnalogueRead > _analogueRefresh) { // Delay for analogue read refresh
         // Issue new read for analogue input states
         _readCommandBuffer[0] = EXIORDAN;
         I2CManager.read(_I2CAddress, _analogueInputBuffer,
@@ -362,14 +382,14 @@ class EXIOExpander : public IODevice {
   uint8_t _minorVer = 0;
   uint8_t _patchVer = 0;
 
-  uint8_t* _digitalInputStates;
-  uint8_t* _analogueInputStates;
-  uint8_t* _analogueInputBuffer;  // buffer for I2C input transfers
+  uint8_t* _digitalInputStates  = NULL;
+  uint8_t* _analogueInputStates = NULL;
+  uint8_t* _analogueInputBuffer = NULL;  // buffer for I2C input transfers
   uint8_t _readCommandBuffer[1];
 
-  uint8_t _digitalPinBytes = 0;  // Size of allocated memory buffer (may be longer than needed)
-  uint8_t _analoguePinBytes = 0;  // Size of allocated memory buffers (may be longer than needed)
-  uint8_t* _analoguePinMap;
+  uint8_t _digitalPinBytes = 0;   // Size of allocated memory buffer (may be longer than needed)
+  uint8_t _analoguePinBytes = 0;  // Size of allocated memory buffer (may be longer than needed)
+  uint8_t* _analoguePinMap = NULL;
   I2CRB _i2crb;
 
   enum {RDS_IDLE, RDS_DIGITAL, RDS_ANALOGUE};  // Read operation states