Skip to content

Commit cdc98bf

Browse files
benrr101benrr101
committed
Renaming the result set reading methods and applying a bit more cleanup on the code within.
1 parent dfe729b commit cdc98bf

1 file changed

Lines changed: 155 additions & 152 deletions

File tree

‎src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs‎

Lines changed: 155 additions & 152 deletions
Original file line numberDiff line numberDiff line change
@@ -780,135 +780,55 @@ private void PrepareTransparentEncryptionFinallyBlock(
780780
}
781781
}
782782

783-
/// <summary>
784-
/// Read the output of sp_describe_parameter_encryption
785-
/// </summary>
786-
/// <param name="ds">Resultset from calling to sp_describe_parameter_encryption</param>
787-
/// <param name="describeParameterEncryptionRpcOriginalRpcMap"> Readonly dictionary with the map of parameter encryption rpc requests with the corresponding original rpc requests.</param>
788-
/// <param name="isRetry">Indicates if this is a retry from a failed call.</param>
789-
private void ReadDescribeEncryptionParameterResults(
790-
SqlDataReader ds, // @TODO: Rename something more obvious
791-
ReadOnlyDictionary<_SqlRPC, _SqlRPC> describeParameterEncryptionRpcOriginalRpcMap,
792-
bool isRetry)
783+
private void ReadDescribeEncryptionParameterAttestation(SqlDataReader ds, bool isRetry)
793784
{
794-
// @TODO: This should be SqlTceCipherInfoTable
795-
Dictionary<int, SqlTceCipherInfoEntry> columnEncryptionKeyTable = new Dictionary<int, SqlTceCipherInfoEntry>();
796-
797-
Debug.Assert((describeParameterEncryptionRpcOriginalRpcMap != null) == _batchRPCMode,
798-
"describeParameterEncryptionRpcOriginalRpcMap should be non-null if and only if it is _batchRPCMode.");
799-
800-
// Indicates the current result set we are reading, used in BatchRPCMode, where we can have more than 1 result set.
801-
int resultSetSequenceNumber = 0;
802-
803-
// A flag that used in BatchRPCMode, to assert the result of lookup in to the dictionary maintaining the map of describe parameter encryption requests
804-
// and the corresponding original rpc requests.
805-
bool lookupDictionaryResult;
806-
807-
// @TODO: If this is supposed to read the results of sp_describe_parameter_encryption there should only ever be 2/3 result sets. So no need to loop this.
808-
do
785+
bool attestationInfoRead = false;
786+
while (ds.Read())
809787
{
810-
if (_batchRPCMode)
811-
{
812-
// If we got more RPC results from the server than what was requested.
813-
if (resultSetSequenceNumber >= _sqlRPCParameterEncryptionReqArray.Length)
814-
{
815-
Debug.Assert(false, "Server sent back more results than what was expected for describe parameter encryption requests in _batchRPCMode.");
816-
// Ignore the rest of the results from the server, if for whatever reason it sends back more than what we expect.
817-
break;
818-
}
819-
}
820-
821-
// 1) Read the first result set that contains the column encryption key list
822-
bool enclaveMetadataExists = ReadDescribeEncryptionParameterResults1(ds, columnEncryptionKeyTable);
823-
if (!enclaveMetadataExists && !ds.NextResult())
824-
{
825-
throw SQL.UnexpectedDescribeParamFormatParameterMetadata();
826-
}
827-
828-
// 2) Find the RPC command that generated this TCE request
829-
_SqlRPC rpc;
830-
if (_batchRPCMode)
831-
{
832-
Debug.Assert(_sqlRPCParameterEncryptionReqArray[resultSetSequenceNumber] != null, "_sqlRPCParameterEncryptionReqArray[resultSetSequenceNumber] should not be null.");
833-
834-
// Lookup in the dictionary to get the original rpc request corresponding to the describe parameter encryption request
835-
// pointed to by _sqlRPCParameterEncryptionReqArray[resultSetSequenceNumber]
836-
rpc = null;
837-
lookupDictionaryResult = describeParameterEncryptionRpcOriginalRpcMap.TryGetValue(_sqlRPCParameterEncryptionReqArray[resultSetSequenceNumber++], out rpc);
838-
839-
Debug.Assert(lookupDictionaryResult,
840-
"Describe Parameter Encryption RPC request key must be present in the dictionary describeParameterEncryptionRpcOriginalRpcMap");
841-
Debug.Assert(rpc != null,
842-
"Describe Parameter Encryption RPC request's corresponding original rpc request must not be null in the dictionary describeParameterEncryptionRpcOriginalRpcMap");
843-
}
844-
else
845-
{
846-
rpc = _rpcArrayOf1[0];
847-
}
848-
849-
Debug.Assert(rpc != null, "rpc should not be null here.");
850-
851-
// 3) Read the second result set containing the cipher metadata
852-
int receivedMetadataCount = 0;
853-
if (!enclaveMetadataExists || ds.NextResult())
854-
{
855-
receivedMetadataCount = ReadDescribeEncryptionParameterResults2(ds, rpc, columnEncryptionKeyTable);
856-
}
857-
858-
// When the RPC object gets reused, the parameter array has more parameters that the valid params for the command.
859-
// Null is used to indicate the end of the valid part of the array. Refer to GetRPCObject().
860-
int userParamCount = rpc.userParams?.Count ?? 0;
861-
if (receivedMetadataCount != userParamCount)
788+
if (attestationInfoRead)
862789
{
863-
for (int index = 0; index < userParamCount; index++)
864-
{
865-
SqlParameter sqlParameter = rpc.userParams[index];
866-
if (!sqlParameter.HasReceivedMetadata && sqlParameter.Direction != ParameterDirection.ReturnValue)
867-
{
868-
// Encryption MD wasn't sent by the server - we expect the metadata to be sent for all the parameters
869-
// that were sent in the original sp_describe_parameter_encryption but not necessarily for return values,
870-
// since there might be multiple return values but server will only send for one of them.
871-
// For parameters that don't need encryption, the encryption type is set to plaintext.
872-
throw SQL.ParamEncryptionMetadataMissing(sqlParameter.ParameterName, rpc.GetCommandTextOrRpcName());
873-
}
874-
}
790+
throw SQL.MultipleRowsReturnedForAttestationInfo();
875791
}
876792

877-
// 4) Read the third result set containing enclave attestation information
878-
if (ShouldUseEnclaveBasedWorkflow && (enclaveAttestationParameters != null) && requiresEnclaveComputations)
879-
{
880-
if (!ds.NextResult())
881-
{
882-
throw SQL.UnexpectedDescribeParamFormatAttestationInfo(this._activeConnection.Parser.EnclaveType);
883-
}
884-
885-
ReadDescribeEncryptionParameterResults3(ds, isRetry);
886-
}
793+
int attestationInfoLength = (int)ds.GetBytes(
794+
(int)DescribeParameterEncryptionResultSet3.AttestationInfo,
795+
dataIndex: 0,
796+
buffer: null,
797+
bufferIndex: 0,
798+
length: 0);
799+
byte[] attestationInfo = new byte[attestationInfoLength];
800+
ds.GetBytes(
801+
(int)DescribeParameterEncryptionResultSet3.AttestationInfo,
802+
dataIndex: 0,
803+
buffer: attestationInfo,
804+
bufferIndex: 0,
805+
length: attestationInfoLength);
887806

888-
// The server has responded with encryption related information for this rpc request. So clear the needsFetchParameterEncryptionMetadata flag.
889-
rpc.needsFetchParameterEncryptionMetadata = false;
890-
} while (ds.NextResult());
807+
SqlConnectionAttestationProtocol attestationProtocol = _activeConnection.AttestationProtocol;
808+
string enclaveType = _activeConnection.Parser.EnclaveType;
891809

892-
// Verify that we received response for each rpc call needs tce
893-
if (_batchRPCMode)
894-
{
895-
for (int i = 0; i < _RPCList.Count; i++)
896-
{
897-
if (_RPCList[i].needsFetchParameterEncryptionMetadata)
898-
{
899-
throw SQL.ProcEncryptionMetadataMissing(_RPCList[i].rpcName);
900-
}
901-
}
810+
EnclaveDelegate.Instance.CreateEnclaveSession(
811+
attestationProtocol,
812+
enclaveType,
813+
GetEnclaveSessionParameters(),
814+
attestationInfo,
815+
enclaveAttestationParameters,
816+
customData,
817+
customDataLength,
818+
isRetry);
819+
enclaveAttestationParameters = null;
820+
attestationInfoRead = true;
902821
}
903822

904-
// If we are not in Batch RPC mode, update the query cache with the encryption MD.
905-
if (!_batchRPCMode && ShouldCacheEncryptionMetadata && (_parameters is not null && _parameters.Count > 0))
823+
if (!attestationInfoRead)
906824
{
907-
SqlQueryMetadataCache.GetInstance().AddQueryMetadata(this, ignoreQueriesWithReturnValueParams: true);
825+
throw SQL.AttestationInfoNotReturnedFromSqlServer(
826+
_activeConnection.Parser.EnclaveType,
827+
_activeConnection.EnclaveAttestationUrl);
908828
}
909829
}
910830

911-
private bool ReadDescribeEncryptionParameterResults1(
831+
private bool ReadDescribeEncryptionParameterKeys(
912832
SqlDataReader ds,
913833
Dictionary<int, SqlTceCipherInfoEntry> columnEncryptionKeyTable)
914834
{
@@ -1048,7 +968,7 @@ private bool ReadDescribeEncryptionParameterResults1(
1048968
return enclaveMetadataExists;
1049969
}
1050970

1051-
private int ReadDescribeEncryptionParameterResults2(
971+
private int ReadDescribeEncryptionParameterMetadata(
1052972
SqlDataReader ds,
1053973
_SqlRPC rpc,
1054974
Dictionary<int, SqlTceCipherInfoEntry> columnEncryptionKeyTable)
@@ -1126,51 +1046,134 @@ private int ReadDescribeEncryptionParameterResults2(
11261046
return receivedMetadataCount;
11271047
}
11281048

1129-
private void ReadDescribeEncryptionParameterResults3(SqlDataReader ds, bool isRetry)
1049+
/// <summary>
1050+
/// Read the output of sp_describe_parameter_encryption
1051+
/// </summary>
1052+
/// <param name="ds">Resultset from calling to sp_describe_parameter_encryption</param>
1053+
/// <param name="describeParameterEncryptionRpcOriginalRpcMap"> Readonly dictionary with the map of parameter encryption rpc requests with the corresponding original rpc requests.</param>
1054+
/// <param name="isRetry">Indicates if this is a retry from a failed call.</param>
1055+
private void ReadDescribeEncryptionParameterResults(
1056+
SqlDataReader ds, // @TODO: Rename something more obvious
1057+
ReadOnlyDictionary<_SqlRPC, _SqlRPC> describeParameterEncryptionRpcOriginalRpcMap,
1058+
bool isRetry)
11301059
{
1131-
bool attestationInfoRead = false;
1132-
while (ds.Read())
1060+
// @TODO: This should be SqlTceCipherInfoTable
1061+
Dictionary<int, SqlTceCipherInfoEntry> columnEncryptionKeyTable = new Dictionary<int, SqlTceCipherInfoEntry>();
1062+
1063+
Debug.Assert(describeParameterEncryptionRpcOriginalRpcMap != null == _batchRPCMode,
1064+
"describeParameterEncryptionRpcOriginalRpcMap should be non-null if and only if it is _batchRPCMode.");
1065+
1066+
// Indicates the current result set we are reading, used in BatchRPCMode, where we can
1067+
// have more than 1 result set.
1068+
int resultSetSequenceNumber = 0;
1069+
1070+
// A flag that used in BatchRPCMode, to assert the result of lookup in to the
1071+
// dictionary maintaining the map of describe parameter encryption requests and the
1072+
// corresponding original rpc requests.
1073+
bool lookupDictionaryResult;
1074+
1075+
// @TODO: If this is supposed to read the results of sp_describe_parameter_encryption there should only ever be 2/3 result sets. So no need to loop this.
1076+
do
11331077
{
1134-
if (attestationInfoRead)
1078+
if (_batchRPCMode)
11351079
{
1136-
throw SQL.MultipleRowsReturnedForAttestationInfo();
1080+
// If we got more RPC results from the server than what was requested.
1081+
if (resultSetSequenceNumber >= _sqlRPCParameterEncryptionReqArray.Length)
1082+
{
1083+
Debug.Fail("Server sent back more results than what was expected for describe parameter encryption requests in _batchRPCMode.");
1084+
// Ignore the rest of the results from the server, if for whatever reason it sends back more than what we expect.
1085+
break;
1086+
}
11371087
}
11381088

1139-
int attestationInfoLength = (int)ds.GetBytes(
1140-
(int)DescribeParameterEncryptionResultSet3.AttestationInfo,
1141-
dataIndex: 0,
1142-
buffer: null,
1143-
bufferIndex: 0,
1144-
length: 0);
1145-
byte[] attestationInfo = new byte[attestationInfoLength];
1146-
ds.GetBytes(
1147-
(int)DescribeParameterEncryptionResultSet3.AttestationInfo,
1148-
dataIndex: 0,
1149-
buffer: attestationInfo,
1150-
bufferIndex: 0,
1151-
length: attestationInfoLength);
1089+
// 1) Read the first result set that contains the column encryption key list
1090+
bool enclaveMetadataExists = ReadDescribeEncryptionParameterKeys(ds, columnEncryptionKeyTable);
1091+
if (!enclaveMetadataExists && !ds.NextResult())
1092+
{
1093+
throw SQL.UnexpectedDescribeParamFormatParameterMetadata();
1094+
}
11521095

1153-
SqlConnectionAttestationProtocol attestationProtocol = _activeConnection.AttestationProtocol;
1154-
string enclaveType = _activeConnection.Parser.EnclaveType;
1096+
// 2) Find the RPC command that generated this TCE request
1097+
_SqlRPC rpc;
1098+
if (_batchRPCMode)
1099+
{
1100+
Debug.Assert(_sqlRPCParameterEncryptionReqArray[resultSetSequenceNumber] != null, "_sqlRPCParameterEncryptionReqArray[resultSetSequenceNumber] should not be null.");
11551101

1156-
EnclaveDelegate.Instance.CreateEnclaveSession(
1157-
attestationProtocol,
1158-
enclaveType,
1159-
GetEnclaveSessionParameters(),
1160-
attestationInfo,
1161-
enclaveAttestationParameters,
1162-
customData,
1163-
customDataLength,
1164-
isRetry);
1165-
enclaveAttestationParameters = null;
1166-
attestationInfoRead = true;
1102+
// Lookup in the dictionary to get the original rpc request corresponding to the describe parameter encryption request
1103+
// pointed to by _sqlRPCParameterEncryptionReqArray[resultSetSequenceNumber]
1104+
lookupDictionaryResult = describeParameterEncryptionRpcOriginalRpcMap.TryGetValue(
1105+
_sqlRPCParameterEncryptionReqArray[resultSetSequenceNumber++],
1106+
out rpc);
1107+
1108+
Debug.Assert(lookupDictionaryResult,
1109+
"Describe Parameter Encryption RPC request key must be present in the dictionary describeParameterEncryptionRpcOriginalRpcMap");
1110+
Debug.Assert(rpc != null,
1111+
"Describe Parameter Encryption RPC request's corresponding original rpc request must not be null in the dictionary describeParameterEncryptionRpcOriginalRpcMap");
1112+
}
1113+
else
1114+
{
1115+
rpc = _rpcArrayOf1[0];
1116+
}
1117+
1118+
Debug.Assert(rpc is not null, "rpc should not be null here.");
1119+
1120+
// 3) Read the second result set containing the per-parameter cipher metadata
1121+
int receivedMetadataCount = 0;
1122+
if (!enclaveMetadataExists || ds.NextResult())
1123+
{
1124+
receivedMetadataCount = ReadDescribeEncryptionParameterMetadata(ds, rpc, columnEncryptionKeyTable);
1125+
}
1126+
1127+
// When the RPC object gets reused, the parameter array has more parameters that the valid params for the command.
1128+
// Null is used to indicate the end of the valid part of the array. Refer to GetRPCObject().
1129+
int userParamCount = rpc.userParams?.Count ?? 0;
1130+
if (receivedMetadataCount != userParamCount)
1131+
{
1132+
for (int index = 0; index < userParamCount; index++)
1133+
{
1134+
SqlParameter sqlParameter = rpc.userParams[index];
1135+
if (!sqlParameter.HasReceivedMetadata && sqlParameter.Direction != ParameterDirection.ReturnValue)
1136+
{
1137+
// Encryption MD wasn't sent by the server - we expect the metadata to be sent for all the parameters
1138+
// that were sent in the original sp_describe_parameter_encryption but not necessarily for return values,
1139+
// since there might be multiple return values but server will only send for one of them.
1140+
// For parameters that don't need encryption, the encryption type is set to plaintext.
1141+
throw SQL.ParamEncryptionMetadataMissing(sqlParameter.ParameterName, rpc.GetCommandTextOrRpcName());
1142+
}
1143+
}
1144+
}
1145+
1146+
// 4) Read the third result set containing enclave attestation information
1147+
if (ShouldUseEnclaveBasedWorkflow && enclaveAttestationParameters != null && requiresEnclaveComputations)
1148+
{
1149+
if (!ds.NextResult())
1150+
{
1151+
throw SQL.UnexpectedDescribeParamFormatAttestationInfo(_activeConnection.Parser.EnclaveType);
1152+
}
1153+
1154+
ReadDescribeEncryptionParameterAttestation(ds, isRetry);
1155+
}
1156+
1157+
// The server has responded with encryption related information for this rpc request. So clear the needsFetchParameterEncryptionMetadata flag.
1158+
rpc.needsFetchParameterEncryptionMetadata = false;
1159+
} while (ds.NextResult());
1160+
1161+
// Verify that we received response for each rpc call needs tce
1162+
if (_batchRPCMode)
1163+
{
1164+
for (int i = 0; i < _RPCList.Count; i++)
1165+
{
1166+
if (_RPCList[i].needsFetchParameterEncryptionMetadata)
1167+
{
1168+
throw SQL.ProcEncryptionMetadataMissing(_RPCList[i].rpcName);
1169+
}
1170+
}
11671171
}
11681172

1169-
if (!attestationInfoRead)
1173+
// If we are not in Batch RPC mode, update the query cache with the encryption MD.
1174+
if (!_batchRPCMode && ShouldCacheEncryptionMetadata && _parameters?.Count > 0)
11701175
{
1171-
throw SQL.AttestationInfoNotReturnedFromSqlServer(
1172-
_activeConnection.Parser.EnclaveType,
1173-
_activeConnection.EnclaveAttestationUrl);
1176+
SqlQueryMetadataCache.GetInstance().AddQueryMetadata(this, ignoreQueriesWithReturnValueParams: true);
11741177
}
11751178
}
11761179

0 commit comments

Comments
 (0)