From b416e166a3eba1dcfc2e912e37e443981d7f60ef Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Tue, 14 Aug 2018 19:24:37 -0700
Subject: [PATCH] rdbms: add domain sanity checks to LoadBalancer connection
 methods

Bug: T193565
Change-Id: I94d905277e01b8e30ac3f6532ece07388bb20cce
---
 .../libs/rdbms/loadbalancer/LoadBalancer.php  | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index d84ba65e3d..ff762b42cc 100644
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -28,6 +28,7 @@ use BagOStuff;
 use EmptyBagOStuff;
 use WANObjectCache;
 use ArrayUtils;
+use UnexpectedValueException;
 use InvalidArgumentException;
 use RuntimeException;
 use Exception;
@@ -952,6 +953,17 @@ class LoadBalancer implements ILoadBalancer {
 			}
 		}
 
+		// Final sanity check to make sure the right domain is selected
+		if (
+			$conn instanceof IDatabase &&
+			$this->localDomain->getDatabase() !== null &&
+			$conn->getDomainID() !== $this->localDomain->getId()
+		) {
+			throw new UnexpectedValueException(
+				"Got connection to '{$conn->getDomainID()}', " .
+				"but expected local domain ('{$this->localDomain}')." );
+		}
+
 		return $conn;
 	}
 
@@ -1038,8 +1050,13 @@ class LoadBalancer implements ILoadBalancer {
 			}
 		}
 
-		// Increment reference count
 		if ( $conn instanceof IDatabase ) {
+			// Final sanity check to make sure the right domain is selected
+			if ( $domainInstance->getDatabase() !== null && $conn->getDomainID() !== $domain ) {
+				throw new UnexpectedValueException(
+					"Got connection to '{$conn->getDomainID()}', but expected '$domain'." );
+			}
+			// Increment reference count
 			$refCount = $conn->getLBInfo( 'foreignPoolRefCount' );
 			$conn->setLBInfo( 'foreignPoolRefCount', $refCount + 1 );
 		}
-- 
2.20.1