unique_ptr<IOBuf> HeaderServerChannel::handleSecurityMessage(
unique_ptr<IOBuf>&& buf) {
if (header_->getClientType() == THRIFT_HEADER_SASL_CLIENT_TYPE) {
if (!header_->isSupportedClient()) {
if (protectionState_ == ProtectionState::UNKNOWN) {
// The client tried to use SASL, but it's not supported by
// policy. Tell the client to fall back.
// TODO mhorowitz: generate a real message here.
try {
sendMessage(nullptr, IOBuf::create(0));
} catch (const std::exception& e) {
LOG(ERROR) << "Failed to send message: " << e.what();
}
return nullptr;
} else {
// The supported client set changed halfway through or
// something. Bail out.
protectionState_ = ProtectionState::INVALID;
LOG(WARNING) << "Inconsistent SASL support";
TTransportException ex("Inconsistent SASL support");
messageReceiveError(make_exception_ptr(ex));
return nullptr;
}
} else if (protectionState_ == ProtectionState::UNKNOWN ||
protectionState_ == ProtectionState::INPROGRESS) {
protectionState_ = ProtectionState::INPROGRESS;
saslServer_->consumeFromClient(&saslServerCallback_, std::move(buf));
return nullptr;
}
// else, fall through to application message processing
} else if (protectionState_ == ProtectionState::VALID ||
(protectionState_ == ProtectionState::INPROGRESS &&
!header_->isSupportedClient())) {
// Either negotiation has completed or negotiation is incomplete,
// non-sasl was received, but is not permitted.
// We should fail hard in this case.
protectionState_ = ProtectionState::INVALID;
LOG(WARNING) << "non-SASL message received on SASL channel";
TTransportException ex("non-SASL message received on SASL channel");
messageReceiveError(make_exception_ptr(ex));
return nullptr;
} else if (protectionState_ == ProtectionState::UNKNOWN) {
// This is the path non-SASL-aware (or SASL-disabled) clients will
// take.
VLOG(5) << "non-SASL client connection received";
protectionState_ = ProtectionState::NONE;
} else if (protectionState_ == ProtectionState::INPROGRESS &&
header_->isSupportedClient()) {
// If a client permits a non-secure connection, we allow falling back to
// one even if a SASL handshake is in progress.
LOG(INFO) << "Client initiated a fallback during a SASL handshake";
// Cancel any SASL-related state, and log
protectionState_ = ProtectionState::NONE;
saslServerCallback_.cancelTimeout();
saslServer_->markChannelCallbackUnavailable();
const auto& observer = getEventBase()->getObserver();
if (observer) {
observer->saslFallBack();
}
}
return std::move(buf);
}
void RequestChannel::sendRequestAsync(
apache::thrift::RpcOptions& rpcOptions,
std::unique_ptr<apache::thrift::RequestCallback> callback,
std::unique_ptr<apache::thrift::ContextStack> ctx,
std::unique_ptr<folly::IOBuf> buf,
std::shared_ptr<apache::thrift::transport::THeader> header,
RpcKind kind) {
auto eb = getEventBase();
if (!eb || eb->isInEventBaseThread()) {
switch (kind) {
case RpcKind::SINGLE_REQUEST_NO_RESPONSE:
// Calling asyncComplete before sending because
// sendOnewayRequest moves from ctx and clears it.
ctx->asyncComplete();
sendOnewayRequest(
rpcOptions,
std::move(callback),
std::move(ctx),
std::move(buf),
std::move(header));
break;
case RpcKind::SINGLE_REQUEST_SINGLE_RESPONSE:
sendRequest(
rpcOptions,
std::move(callback),
std::move(ctx),
std::move(buf),
std::move(header));
break;
case RpcKind::SINGLE_REQUEST_STREAMING_RESPONSE:
sendStreamRequest(
rpcOptions,
std::move(callback),
std::move(ctx),
std::move(buf),
std::move(header));
break;
default:
folly::assume_unreachable();
break;
}
} else {
switch (kind) {
case RpcKind::SINGLE_REQUEST_NO_RESPONSE:
eb->runInEventBaseThread([this,
rpcOptions,
callback = std::move(callback),
ctx = std::move(ctx),
buf = std::move(buf),
header = std::move(header)]() mutable {
// Calling asyncComplete before sending because
// sendOnewayRequest moves from ctx and clears it.
ctx->asyncComplete();
sendOnewayRequest(
rpcOptions,
std::move(callback),
std::move(ctx),
std::move(buf),
std::move(header));
});
break;
case RpcKind::SINGLE_REQUEST_SINGLE_RESPONSE:
eb->runInEventBaseThread([this,
rpcOptions,
callback = std::move(callback),
ctx = std::move(ctx),
buf = std::move(buf),
header = std::move(header)]() mutable {
sendRequest(
rpcOptions,
std::move(callback),
std::move(ctx),
std::move(buf),
std::move(header));
});
break;
case RpcKind::SINGLE_REQUEST_STREAMING_RESPONSE:
eb->runInEventBaseThread([this,
rpcOptions,
callback = std::move(callback),
ctx = std::move(ctx),
buf = std::move(buf),
header = std::move(header)]() mutable {
sendStreamRequest(
rpcOptions,
std::move(callback),
std::move(ctx),
std::move(buf),
std::move(header));
});
break;
default:
folly::assume_unreachable();
break;
}
}
}
testing::NiceMock<MockTAsyncTransport>* newMockTransport(EventBase* evb) {
auto transport = new testing::NiceMock<MockTAsyncTransport>();
EXPECT_CALL(*transport, getEventBase())
.WillRepeatedly(testing::Return(evb));
return transport;
}
void RequestChannel::sendRequestSync(
RpcOptions& options,
std::unique_ptr<RequestCallback> cb,
std::unique_ptr<apache::thrift::ContextStack> ctx,
std::unique_ptr<folly::IOBuf> buf,
std::shared_ptr<apache::thrift::transport::THeader> header,
RpcKind kind) {
auto eb = getEventBase();
// We intentionally only support sync_* calls from the EventBase thread.
eb->checkIsInEventBaseThread();
folly::fibers::Baton baton;
auto scb =
std::make_unique<ClientSyncBatonCallback>(std::move(cb), baton, kind);
folly::exception_wrapper ew;
auto const onFiber = folly::fibers::onFiber();
auto const inEventBase = eb->inRunningEventBaseThread();
baton.wait([&, onFiber, inEventBase] {
try {
switch (kind) {
case RpcKind::SINGLE_REQUEST_NO_RESPONSE:
sendOnewayRequest(
options,
std::move(scb),
std::move(ctx),
std::move(buf),
std::move(header));
break;
case RpcKind::SINGLE_REQUEST_SINGLE_RESPONSE:
sendRequest(
options,
std::move(scb),
std::move(ctx),
std::move(buf),
std::move(header));
break;
case RpcKind::SINGLE_REQUEST_STREAMING_RESPONSE:
sendStreamRequest(
options,
std::move(scb),
std::move(ctx),
std::move(buf),
std::move(header));
break;
default:
folly::assume_unreachable();
}
} catch (const std::exception& e) {
ew = folly::exception_wrapper(std::current_exception(), e);
baton.post();
} catch (...) {
ew = folly::exception_wrapper(std::current_exception());
baton.post();
}
if (!onFiber || !inEventBase) {
while (!baton.ready()) {
eb->drive();
}
}
});
if (ew) {
ew.throw_exception();
}
}