schanneladapter.cc

本人收集整理的一份c/c++跨平台网络库

      if ((data_len == 0) && (inbuf[0] == 0x15)) {
        status = SEC_I_CONTEXT_EXPIRED;
      }
      if (extra_len) {
        size_t consumed = inbuf.size() - extra_len;
        memmove(&inbuf[0], &inbuf[consumed], extra_len);
        inbuf.resize(extra_len);
      } else {
        inbuf.clear();
      }
      // TODO: Handle SEC_I_CONTEXT_EXPIRED to do clean shutdown
      if (status != SEC_E_OK) {
        LOG(LS_INFO) << "DecryptMessage returned continuation code: "
                      << ErrorName(status, SECURITY_ERRORS);
      }
      continue;
    }
    
    if (status == SEC_E_INCOMPLETE_MESSAGE) {
      break;
    } else {
      return status;
    }
  }

  return 0;
}

void
SChannelAdapter::Cleanup() {
  if (impl_->ctx_init)
    DeleteSecurityContext(&impl_->ctx);
  if (impl_->cred_init)
    FreeCredentialsHandle(&impl_->cred);
  delete impl_;
}

void
SChannelAdapter::PostEvent() {
  // Check if there's anything notable to signal
  if (impl_->readable.empty() && !signal_close_)
    return;

  // Only one post in the queue at a time
  if (message_pending_)
    return;

  if (Thread* thread = Thread::Current()) {
    message_pending_ = true;
    thread->Post(this);
  } else {
    LOG(LS_ERROR) << "No thread context available for SChannelAdapter";
    ASSERT(false);
  }
}

void
SChannelAdapter::Error(const char* context, int err, bool signal) {
  LOG(LS_WARNING) << "SChannelAdapter::Error("
                  << context << ", "
                  << ErrorName(err, SECURITY_ERRORS) << ")";
  state_ = SSL_ERROR;
  SetError(err);
  if (signal)
    AsyncSocketAdapter::OnCloseEvent(this, err);
}

int
SChannelAdapter::Read() {
  char buffer[4096];
  SChannelBuffer& inbuf = impl_->inbuf;
  while (true) {
    int ret = AsyncSocketAdapter::Recv(buffer, sizeof(buffer));
    if (ret > 0) {
      inbuf.insert(inbuf.end(), buffer, buffer + ret);
    } else if (GetError() == EWOULDBLOCK) {
      return 0;  // Blocking
    } else {
      return GetError();
    }
  }
}

int
SChannelAdapter::Flush() {
  int result = 0;
  size_t pos = 0;
  SChannelBuffer& outbuf = impl_->outbuf;
  while (pos < outbuf.size()) {
    int sent = AsyncSocketAdapter::Send(&outbuf[pos], outbuf.size() - pos);
    if (sent > 0) {
      pos += sent;
    } else if (GetError() == EWOULDBLOCK) {
      break;  // Blocking
    } else {
      result = GetError();
      break;
    }
  }
  if (int remainder = (int)(outbuf.size() - pos)) {
    memmove(&outbuf[0], &outbuf[pos], remainder);
    outbuf.resize(remainder);
  } else {
    outbuf.clear();
  }
  return result;
}

//
// AsyncSocket Implementation
//

int
SChannelAdapter::Send(const void* pv, size_t cb) {
  switch (state_) {
  case SSL_NONE:
    return AsyncSocketAdapter::Send(pv, cb);

  case SSL_WAIT:
  case SSL_CONNECTING:
    SetError(EWOULDBLOCK);
    return SOCKET_ERROR;

  case SSL_CONNECTED:
    break;

  case SSL_ERROR:
  default:
    return SOCKET_ERROR;
  }

  size_t written = 0;
  SChannelBuffer& outbuf = impl_->outbuf;
  while (written < cb) {
    const size_t encrypt_len = std::min<size_t>(cb - written,
                                                impl_->sizes.cbMaximumMessage);

    CSecBufferBundle<4> out_buf;
    out_buf[0].BufferType = SECBUFFER_STREAM_HEADER;
    out_buf[0].cbBuffer = impl_->sizes.cbHeader;
    out_buf[1].BufferType = SECBUFFER_DATA;
    out_buf[1].cbBuffer = static_cast<unsigned long>(encrypt_len);
    out_buf[2].BufferType = SECBUFFER_STREAM_TRAILER;
    out_buf[2].cbBuffer = impl_->sizes.cbTrailer;

    size_t packet_len = out_buf[0].cbBuffer
                      + out_buf[1].cbBuffer
                      + out_buf[2].cbBuffer;

    SChannelBuffer message;
    message.resize(packet_len);
    out_buf[0].pvBuffer = &message[0];
    out_buf[1].pvBuffer = &message[out_buf[0].cbBuffer];
    out_buf[2].pvBuffer = &message[out_buf[0].cbBuffer + out_buf[1].cbBuffer];

    memcpy(out_buf[1].pvBuffer,
           static_cast<const char*>(pv) + written,
           encrypt_len);

    //DescribeBuffers(LS_VERBOSE, "Encrypt In ", out_buf.desc());
    SECURITY_STATUS res = EncryptMessage(&impl_->ctx, 0, out_buf.desc(), 0);
    //DescribeBuffers(LS_VERBOSE, "Encrypt Out ", out_buf.desc());

    if (FAILED(res)) {
      Error("EncryptMessage", res, false);
      return SOCKET_ERROR;
    }

    // We assume that the header and data segments do not change length,
    // or else encrypting the concatenated packet in-place is wrong.
    ASSERT(out_buf[0].cbBuffer == impl_->sizes.cbHeader);
    ASSERT(out_buf[1].cbBuffer == static_cast<unsigned long>(encrypt_len));

    // However, the length of the trailer may change due to padding.
    ASSERT(out_buf[2].cbBuffer <= impl_->sizes.cbTrailer);

    packet_len = out_buf[0].cbBuffer
               + out_buf[1].cbBuffer
               + out_buf[2].cbBuffer;

    written += encrypt_len;
    outbuf.insert(outbuf.end(), &message[0], &message[packet_len-1]+1);
  }

  if (int err = Flush()) {
    state_ = SSL_ERROR;
    SetError(err);
    return SOCKET_ERROR;
  }

  return static_cast<int>(written);
}

int
SChannelAdapter::Recv(void* pv, size_t cb) {
  switch (state_) {
  case SSL_NONE:
    return AsyncSocketAdapter::Recv(pv, cb);

  case SSL_WAIT:
  case SSL_CONNECTING:
    SetError(EWOULDBLOCK);
    return SOCKET_ERROR;

  case SSL_CONNECTED:
    break;

  case SSL_ERROR:
  default:
    return SOCKET_ERROR;
  }

  SChannelBuffer& readable = impl_->readable;
  if (readable.empty()) {
    SetError(EWOULDBLOCK);
    return SOCKET_ERROR;
  }
  size_t read = min(cb, readable.size());
  memcpy(pv, &readable[0], read);
  if (size_t remaining = readable.size() - read) {
    memmove(&readable[0], &readable[read], remaining);
    readable.resize(remaining);
  } else {
    readable.clear();
  }

  PostEvent();
  return static_cast<int>(read);
}

int
SChannelAdapter::Close() {
  if (!impl_->readable.empty()) {
    LOG(WARNING) << "SChannelAdapter::Close with readable data";
    // Note: this isn't strictly an error, but we're using it temporarily to
    // track bugs.
    //ASSERT(false);
  }
  if (state_ == SSL_CONNECTED) {
    DWORD token = SCHANNEL_SHUTDOWN;
    CSecBufferBundle<1> sb_in;
    sb_in[0].BufferType = SECBUFFER_TOKEN;
    sb_in[0].cbBuffer = sizeof(token);
    sb_in[0].pvBuffer = &token;
    ApplyControlToken(&impl_->ctx, sb_in.desc());
    // TODO: In theory, to do a nice shutdown, we need to begin shutdown
    // negotiation with more calls to InitializeSecurityContext.  Since the
    // socket api doesn't support nice shutdown at this point, we don't bother.
  }
  Cleanup();
  impl_ = new SSLImpl;
  state_ = restartable_ ? SSL_WAIT : SSL_NONE;
  signal_close_ = false;
  message_pending_ = false;
  return AsyncSocketAdapter::Close();
}

Socket::ConnState
SChannelAdapter::GetState() const {
  if (signal_close_)
    return CS_CONNECTED;
  ConnState state = socket_->GetState();
  if ((state == CS_CONNECTED)
      && ((state_ == SSL_WAIT) || (state_ == SSL_CONNECTING)))
    state = CS_CONNECTING;
  return state;
}

void
SChannelAdapter::OnConnectEvent(AsyncSocket* socket) {
  LOG(LS_VERBOSE) << "SChannelAdapter::OnConnectEvent";
  if (state_ != SSL_WAIT) {
    ASSERT(state_ == SSL_NONE);
    AsyncSocketAdapter::OnConnectEvent(socket);
    return;
  }

  state_ = SSL_CONNECTING;
  if (int err = BeginSSL()) {
    Error("BeginSSL", err);
  }
}

void
SChannelAdapter::OnReadEvent(AsyncSocket* socket) {
  if (state_ == SSL_NONE) {
    AsyncSocketAdapter::OnReadEvent(socket);
    return;
  }

  if (int err = Read()) {
    Error("Read", err);
    return;
  }

  if (impl_->inbuf.empty())
    return;

  if (state_ == SSL_CONNECTED) {
    if (int err = DecryptData()) {
      Error("DecryptData", err);
    } else if (!impl_->readable.empty()) {
      AsyncSocketAdapter::OnReadEvent(this);
    }
  } else if (state_ == SSL_CONNECTING) {
    if (int err = ContinueSSL()) {
      Error("ContinueSSL", err);
    }
  }
}

void
SChannelAdapter::OnWriteEvent(AsyncSocket* socket) {
  if (state_ == SSL_NONE) {
    AsyncSocketAdapter::OnWriteEvent(socket);
    return;
  }

  if (int err = Flush()) {
    Error("Flush", err);
    return;
  }

  // See if we have more data to write
  if (!impl_->outbuf.empty())
    return;

  // Buffer is empty, submit notification
  if (state_ == SSL_CONNECTED) {
    AsyncSocketAdapter::OnWriteEvent(socket);
  }
}

void
SChannelAdapter::OnCloseEvent(AsyncSocket* socket, int err) {
  if ((state_ == SSL_NONE) || impl_->readable.empty()) {
    AsyncSocketAdapter::OnCloseEvent(socket, err);
    return;
  }

  // If readable is non-empty, then we have a pending Message
  // that will allow us to signal close (eventually).
  signal_close_ = true;
}

void
SChannelAdapter::OnMessage(Message* pmsg) {
  if (!message_pending_)
    return;  // This occurs when socket is closed

  message_pending_ = false;
  if (!impl_->readable.empty()) {
    AsyncSocketAdapter::OnReadEvent(this);
  } else if (signal_close_) {
    signal_close_ = false;
    AsyncSocketAdapter::OnCloseEvent(this, 0); // TODO: cache this error?
  }
}

} // namespace utils_base