Implement BIO_write_ex

This does not yet switch BoringSSL's built-in BIOs, and callers within
the library to the new function. That will be done in subsequent CLs.

Bug: 42290376
Change-Id: Id679908a02f04e53487ec68f5931a0dd8ec8e405
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/96109
Commit-Queue: David Benjamin <davidben@google.com>
Reviewed-by: Rudolf Polzer <rpolzer@google.com>
diff --git a/crypto/bio/bio.cc b/crypto/bio/bio.cc
index 793389c..17b068c 100644
--- a/crypto/bio/bio.cc
+++ b/crypto/bio/bio.cc
@@ -19,6 +19,7 @@
 #include <limits.h>
 #include <string.h>
 
+#include <algorithm>
 #include <utility>
 
 #include <openssl/asn1.h>
@@ -119,27 +120,56 @@
   return ret;
 }
 
-int BIO_write(BIO *bio, const void *in, int inl) {
+int BIO_write_ex(BIO *bio, const void *data, size_t len, size_t *out_written) {
+  if (out_written != nullptr) {
+    *out_written = 0;
+  }
   auto *impl = FromOpaque(bio);
-  if (impl == nullptr || impl->method->bwrite == nullptr) {
+  if (impl == nullptr ||
+      (impl->method->bwrite == nullptr && impl->method->bwrite_ex == nullptr)) {
     OPENSSL_PUT_ERROR(BIO, BIO_R_UNSUPPORTED_METHOD);
-    return -1;
+    return 0;
   }
   if (!impl->init) {
     OPENSSL_PUT_ERROR(BIO, BIO_R_UNINITIALIZED);
-    return -1;
-  }
-  if (inl <= 0) {
     return 0;
   }
-  int ret = impl->method->bwrite(impl, reinterpret_cast<const char *>(in), inl);
-  if (ret > 0) {
-    impl->num_write += ret;
-  } else {
-    // In preparation for |BIO_write_ex|, canonicalize error returns to -1.
-    ret = -1;
+  // Matching OpenSSL, writing zero bytes should "successfully" write no bytes.
+  if (len == 0) {
+    return 1;
   }
-  return ret;
+  // Support `BIO_METHOD`s using either the old or new API.
+  size_t written;
+  if (impl->method->bwrite_ex != nullptr) {
+    if (!impl->method->bwrite_ex(bio, static_cast<const char *>(data), len,
+                                 &written)) {
+      return 0;
+    }
+  } else {
+    int ret =
+        impl->method->bwrite(impl, static_cast<const char *>(data),
+                             static_cast<int>(std::min(len, size_t{INT_MAX})));
+    if (ret <= 0) {
+      return 0;
+    }
+    written = static_cast<size_t>(ret);
+  }
+  impl->num_write += written;
+  if (out_written != nullptr) {
+    *out_written = written;
+  }
+  return 1;
+}
+
+int BIO_write(BIO *bio, const void *in, int inl) {
+  // Matching OpenSSL, `inl <= 0` returns zero, i.e. a "successful" write of
+  // zero bytes.
+  inl = std::max(inl, 0);
+  size_t written;
+  if (!BIO_write_ex(bio, in, static_cast<size_t>(inl), &written)) {
+    return -1;
+  }
+  return static_cast<int>(written);
 }
 
 int BIO_write_all(BIO *bio, const void *data, size_t len) {
@@ -613,6 +643,13 @@
   return 1;
 }
 
+int BIO_meth_set_write_ex(BIO_METHOD *method,
+                          int (*write_ex_func)(BIO *, const char *, size_t,
+                                               size_t *)) {
+  method->bwrite_ex = write_ex_func;
+  return 1;
+}
+
 int BIO_meth_set_write(BIO_METHOD *method,
                        int (*write_func)(BIO *, const char *, int)) {
   method->bwrite = write_func;
diff --git a/crypto/bio/bio_mem.cc b/crypto/bio/bio_mem.cc
index 479b9ce..2d80db7 100644
--- a/crypto/bio/bio_mem.cc
+++ b/crypto/bio/bio_mem.cc
@@ -224,9 +224,11 @@
 }
 
 static const BIO_METHOD mem_method = {
-    BIO_TYPE_MEM, "memory buffer", mem_write,
-    mem_read,     mem_gets,        mem_ctrl,
-    mem_new,      mem_free,        /*callback_ctrl=*/nullptr,
+    BIO_TYPE_MEM, "memory buffer",
+    mem_write,    /*bwrite_ex=*/nullptr,
+    mem_read,     mem_gets,
+    mem_ctrl,     mem_new,
+    mem_free,     /*callback_ctrl=*/nullptr,
 };
 
 const BIO_METHOD *BIO_s_mem() { return &mem_method; }
diff --git a/crypto/bio/bio_test.cc b/crypto/bio/bio_test.cc
index faaa87f..eecb81e 100644
--- a/crypto/bio/bio_test.cc
+++ b/crypto/bio/bio_test.cc
@@ -1449,5 +1449,132 @@
   EXPECT_EQ(BIO_write(bio.get(), "foo", 3), -1);
 }
 
+// BIO_write should work with BIOs that implement BIO_write_ex and vice versa.
+TEST(BIOTest, WriteExConversion) {
+  enum class WriteResult { kSuccess, kFail, kRetry };
+  struct MockWrite {
+    std::string expected_data;
+    WriteResult result = WriteResult::kSuccess;
+    size_t bytes_written = 0;
+    int fail_return = -1;  // Only used for write_method failure
+  };
+
+  UniquePtr<BIO_METHOD> write_method(BIO_meth_new(0, nullptr));
+  ASSERT_TRUE(write_method);
+  BIO_meth_set_write(
+      write_method.get(), [](BIO *bio, const char *in, int len) -> int {
+        auto *state = static_cast<MockWrite *>(BIO_get_data(bio));
+        EXPECT_EQ(std::string(in, len), state->expected_data);
+        BIO_clear_retry_flags(bio);
+        switch (state->result) {
+          case WriteResult::kSuccess:
+            return static_cast<int>(state->bytes_written);
+          case WriteResult::kFail:
+            return state->fail_return;
+          case WriteResult::kRetry:
+            BIO_set_retry_write(bio);
+            return -1;
+        }
+        return -1;
+      });
+
+  UniquePtr<BIO_METHOD> write_ex_method(BIO_meth_new(0, nullptr));
+  ASSERT_TRUE(write_ex_method);
+  BIO_meth_set_write_ex(
+      write_ex_method.get(),
+      [](BIO *bio, const char *in, size_t len, size_t *out_written) -> int {
+        auto *state = static_cast<MockWrite *>(BIO_get_data(bio));
+        EXPECT_EQ(std::string(in, len), state->expected_data);
+        BIO_clear_retry_flags(bio);
+        switch (state->result) {
+          case WriteResult::kSuccess:
+            *out_written = state->bytes_written;
+            return 1;
+          case WriteResult::kFail:
+            return 0;
+          case WriteResult::kRetry:
+            BIO_set_retry_write(bio);
+            return 0;
+        }
+        return 0;
+      });
+
+  for (const BIO_METHOD *meth : {write_method.get(), write_ex_method.get()}) {
+    SCOPED_TRACE(meth == write_method.get() ? "write" : "write_ex");
+
+    auto make_bio = [&](MockWrite *mock_write) -> UniquePtr<BIO> {
+      UniquePtr<BIO> bio(BIO_new(meth));
+      if (bio == nullptr) {
+        return nullptr;
+      }
+      BIO_set_data(bio.get(), mock_write);
+      BIO_set_init(bio.get(), 1);
+      return bio;
+    };
+
+    MockWrite mock_write;
+    static const char kInput[] = "hello";
+    mock_write.expected_data = kInput;
+
+    // Test BIO_write
+    {
+      auto bio = make_bio(&mock_write);
+      ASSERT_TRUE(bio);
+
+      mock_write.result = WriteResult::kSuccess;
+      mock_write.bytes_written = 5;
+      EXPECT_EQ(BIO_write(bio.get(), kInput, strlen(kInput)), 5);
+      EXPECT_FALSE(BIO_should_retry(bio.get()));
+
+      mock_write.result = WriteResult::kFail;
+      mock_write.fail_return = -1;
+      EXPECT_EQ(BIO_write(bio.get(), kInput, strlen(kInput)), -1);
+      EXPECT_FALSE(BIO_should_retry(bio.get()));
+
+      // Test both return values for |write_method|.
+      mock_write.fail_return = 0;
+      EXPECT_EQ(BIO_write(bio.get(), kInput, strlen(kInput)), -1);
+      EXPECT_FALSE(BIO_should_retry(bio.get()));
+
+      mock_write.result = WriteResult::kRetry;
+      EXPECT_EQ(BIO_write(bio.get(), kInput, strlen(kInput)), -1);
+      EXPECT_TRUE(BIO_should_retry(bio.get()));
+      EXPECT_TRUE(BIO_should_write(bio.get()));
+    }
+
+    // Test BIO_write_ex
+    {
+      auto bio = make_bio(&mock_write);
+      ASSERT_TRUE(bio);
+      size_t written = 0;
+
+      mock_write.result = WriteResult::kSuccess;
+      mock_write.bytes_written = 5;
+      EXPECT_EQ(BIO_write_ex(bio.get(), kInput, strlen(kInput), &written), 1);
+      EXPECT_EQ(written, 5u);
+      EXPECT_FALSE(BIO_should_retry(bio.get()));
+
+      // `out_written` can be nullptr.
+      EXPECT_EQ(BIO_write_ex(bio.get(), kInput, strlen(kInput), nullptr), 1);
+      EXPECT_FALSE(BIO_should_retry(bio.get()));
+
+      mock_write.result = WriteResult::kFail;
+      mock_write.fail_return = -1;
+      EXPECT_EQ(BIO_write_ex(bio.get(), kInput, strlen(kInput), &written), 0);
+      EXPECT_FALSE(BIO_should_retry(bio.get()));
+
+      // Test both return values for |write_method|.
+      mock_write.fail_return = 0;
+      EXPECT_EQ(BIO_write_ex(bio.get(), kInput, strlen(kInput), &written), 0);
+      EXPECT_FALSE(BIO_should_retry(bio.get()));
+
+      mock_write.result = WriteResult::kRetry;
+      EXPECT_EQ(BIO_write_ex(bio.get(), kInput, strlen(kInput), &written), 0);
+      EXPECT_TRUE(BIO_should_retry(bio.get()));
+      EXPECT_TRUE(BIO_should_write(bio.get()));
+    }
+  }
+}
+
 }  // namespace
 BSSL_NAMESPACE_END
diff --git a/crypto/bio/connect.cc b/crypto/bio/connect.cc
index 328c2b1..d2d478e 100644
--- a/crypto/bio/connect.cc
+++ b/crypto/bio/connect.cc
@@ -425,9 +425,9 @@
 }
 
 static const BIO_METHOD methods_connectp = {
-    BIO_TYPE_CONNECT, "socket connect", conn_write,
-    conn_read,        /*gets=*/nullptr, conn_ctrl,
-    conn_new,         conn_free,        conn_callback_ctrl,
+    BIO_TYPE_CONNECT, "socket connect",   conn_write, /*bwrite_ex=*/nullptr,
+    conn_read,        /*gets=*/nullptr,   conn_ctrl,  conn_new,
+    conn_free,        conn_callback_ctrl,
 };
 
 const BIO_METHOD *BIO_s_connect() { return &methods_connectp; }
diff --git a/crypto/bio/fd.cc b/crypto/bio/fd.cc
index b49acb6..8e6058c 100644
--- a/crypto/bio/fd.cc
+++ b/crypto/bio/fd.cc
@@ -163,9 +163,11 @@
 }
 
 static const BIO_METHOD methods_fdp = {
-    BIO_TYPE_FD, "file descriptor", fd_write,
-    fd_read,     fd_gets,           fd_ctrl,
-    fd_new,      fd_free,           /*callback_ctrl=*/nullptr,
+    BIO_TYPE_FD, "file descriptor",
+    fd_write,    /*bwrite_ex=*/nullptr,
+    fd_read,     fd_gets,
+    fd_ctrl,     fd_new,
+    fd_free,     /*callback_ctrl=*/nullptr,
 };
 
 const BIO_METHOD *BIO_s_fd() { return &methods_fdp; }
diff --git a/crypto/bio/file.cc b/crypto/bio/file.cc
index ca80b69..1481e14 100644
--- a/crypto/bio/file.cc
+++ b/crypto/bio/file.cc
@@ -234,9 +234,16 @@
 }
 
 static const BIO_METHOD methods_filep = {
-    BIO_TYPE_FILE,      "FILE pointer", file_write,
-    file_read,          file_gets,      file_ctrl,
-    /*create=*/nullptr, file_free,      /*callback_ctrl=*/nullptr,
+    BIO_TYPE_FILE,
+    "FILE pointer",
+    file_write,
+    /*bwrite_ex=*/nullptr,
+    file_read,
+    file_gets,
+    file_ctrl,
+    /*create=*/nullptr,
+    file_free,
+    /*callback_ctrl=*/nullptr,
 };
 
 const BIO_METHOD *BIO_s_file() { return &methods_filep; }
diff --git a/crypto/bio/internal.h b/crypto/bio/internal.h
index 72215fc..18a13a3 100644
--- a/crypto/bio/internal.h
+++ b/crypto/bio/internal.h
@@ -43,6 +43,7 @@
   int type;
   const char *name;
   int (*bwrite)(BIO *, const char *, int);
+  int (*bwrite_ex)(BIO *, const char *, size_t, size_t *);
   int (*bread)(BIO *, char *, int);
   int (*bgets)(BIO *, char *, int);
   long (*ctrl)(BIO *, int, long, void *);
diff --git a/crypto/bio/pair.cc b/crypto/bio/pair.cc
index 11b5c2e..828b2a9 100644
--- a/crypto/bio/pair.cc
+++ b/crypto/bio/pair.cc
@@ -402,6 +402,7 @@
     BIO_TYPE_BIO,
     "BIO pair",
     bio_write,
+    /*bwrite_ex=*/nullptr,
     bio_read,
     /*gets=*/nullptr,
     bio_ctrl,
diff --git a/crypto/bio/socket.cc b/crypto/bio/socket.cc
index c61daec..88ede9a 100644
--- a/crypto/bio/socket.cc
+++ b/crypto/bio/socket.cc
@@ -112,15 +112,11 @@
 }
 
 static const BIO_METHOD methods_sockp = {
-    BIO_TYPE_SOCKET,
-    "socket",
-    sock_write,
-    sock_read,
-    nullptr /* gets, */,
-    sock_ctrl,
-    nullptr /* create */,
-    sock_free,
-    nullptr /* callback_ctrl */,
+    BIO_TYPE_SOCKET, "socket",
+    sock_write,      /*bwrite_ex=*/nullptr,
+    sock_read,       nullptr /* gets, */,
+    sock_ctrl,       nullptr /* create */,
+    sock_free,       nullptr /* callback_ctrl */,
 };
 
 const BIO_METHOD *BIO_s_socket() { return &methods_sockp; }
diff --git a/decrepit/bio/base64_bio.cc b/decrepit/bio/base64_bio.cc
index b799999..f0c7a9d 100644
--- a/decrepit/bio/base64_bio.cc
+++ b/decrepit/bio/base64_bio.cc
@@ -480,9 +480,9 @@
 }
 
 static const BIO_METHOD b64_method = {
-    BIO_TYPE_BASE64,   "base64 encoding", b64_write, b64_read,
-    /*bgets=*/nullptr, b64_ctrl,          b64_new,   b64_free,
-    b64_callback_ctrl,
+    BIO_TYPE_BASE64, "base64 encoding", b64_write, /*bwrite_ex=*/nullptr,
+    b64_read,        /*bgets=*/nullptr, b64_ctrl,  b64_new,
+    b64_free,        b64_callback_ctrl,
 };
 
 const BIO_METHOD *BIO_f_base64() { return &b64_method; }
diff --git a/include/openssl/bio.h b/include/openssl/bio.h
index d994c11..a7fcd6a 100644
--- a/include/openssl/bio.h
+++ b/include/openssl/bio.h
@@ -82,6 +82,12 @@
 // outputs the bytes which were available.
 OPENSSL_EXPORT int BIO_gets(BIO *bio, char *buf, int size);
 
+// BIO_write_ex writes `len` bytes from `data` to `bio`. On success, it returns
+// one and sets `*out_written` to the number of bytes written. Otherwise, it
+// returns zero. `out_written` may be NULL to ignore the value.
+OPENSSL_EXPORT int BIO_write_ex(BIO *bio, const void *data, size_t len,
+                                size_t *out_written);
+
 // BIO_write writes `len` bytes from `data` to `bio`. It returns the number of
 // bytes written or a negative number on error.
 OPENSSL_EXPORT int BIO_write(BIO *bio, const void *data, int len);
@@ -682,9 +688,28 @@
 OPENSSL_EXPORT int BIO_meth_set_destroy(BIO_METHOD *method,
                                         int (*destroy_func)(BIO *));
 
+// BIO_meth_set_write_ex sets the implementation of `BIO_write_ex` for `method`
+// and returns one. `BIO_METHOD`s which implement `BIO_write_ex` should also
+// implement `BIO_CTRL_FLUSH`. (See `BIO_meth_set_ctrl`.)
+//
+// `write_ex_func` can assume `out_written` is non-NULL.
+//
+// If configured, `write_ex_func` will also be used to implement `BIO_write`,
+// with the `BIO` framework converting the conventions. `BIO_meth_set_write_ex`
+// and `BIO_meth_set_write` should not be configured on the same `BIO_METHOD`.
+// Prefer `BIO_meth_set_write_ex` for a `size_t`-based API.
+OPENSSL_EXPORT int BIO_meth_set_write_ex(
+    BIO_METHOD *method,
+    int (*write_ex_func)(BIO *, const char *, size_t, size_t *));
+
 // BIO_meth_set_write sets the implementation of `BIO_write` for `method` and
 // returns one. `BIO_METHOD`s which implement `BIO_write` should also implement
 // `BIO_CTRL_FLUSH`. (See `BIO_meth_set_ctrl`.)
+//
+// If configured, `write_func` will also be used to implement `BIO_write_ex`,
+// with the `BIO` framework converting the conventions. `BIO_meth_set_write_ex`
+// and `BIO_meth_set_write` should not be configured on the same `BIO_METHOD`.
+// Prefer `BIO_meth_set_write_ex` for a `size_t`-based API.
 OPENSSL_EXPORT int BIO_meth_set_write(BIO_METHOD *method,
                                       int (*write_func)(BIO *, const char *,
                                                         int));
diff --git a/include/openssl/prefix_symbols.h b/include/openssl/prefix_symbols.h
index 5b4bace..5adae37 100644
--- a/include/openssl/prefix_symbols.h
+++ b/include/openssl/prefix_symbols.h
@@ -274,6 +274,7 @@
 #pragma redefine_extname BIO_meth_set_puts BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_meth_set_puts)
 #pragma redefine_extname BIO_meth_set_read BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_meth_set_read)
 #pragma redefine_extname BIO_meth_set_write BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_meth_set_write)
+#pragma redefine_extname BIO_meth_set_write_ex BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_meth_set_write_ex)
 #pragma redefine_extname BIO_method_type BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_method_type)
 #pragma redefine_extname BIO_new BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_new)
 #pragma redefine_extname BIO_new_bio_pair BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_new_bio_pair)
@@ -337,6 +338,7 @@
 #pragma redefine_extname BIO_wpending BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_wpending)
 #pragma redefine_extname BIO_write BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_write)
 #pragma redefine_extname BIO_write_all BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_write_all)
+#pragma redefine_extname BIO_write_ex BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_write_ex)
 #pragma redefine_extname BIO_write_filename BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BIO_write_filename)
 #pragma redefine_extname BLAKE2B256 BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BLAKE2B256)
 #pragma redefine_extname BLAKE2B256_Final BORINGSSL_ADD_USER_LABEL_AND_PREFIX(BLAKE2B256_Final)
@@ -3388,6 +3390,7 @@
 #define BIO_meth_set_puts BORINGSSL_ADD_PREFIX(BIO_meth_set_puts)
 #define BIO_meth_set_read BORINGSSL_ADD_PREFIX(BIO_meth_set_read)
 #define BIO_meth_set_write BORINGSSL_ADD_PREFIX(BIO_meth_set_write)
+#define BIO_meth_set_write_ex BORINGSSL_ADD_PREFIX(BIO_meth_set_write_ex)
 #define BIO_method_type BORINGSSL_ADD_PREFIX(BIO_method_type)
 #define BIO_new BORINGSSL_ADD_PREFIX(BIO_new)
 #define BIO_new_bio_pair BORINGSSL_ADD_PREFIX(BIO_new_bio_pair)
@@ -3451,6 +3454,7 @@
 #define BIO_wpending BORINGSSL_ADD_PREFIX(BIO_wpending)
 #define BIO_write BORINGSSL_ADD_PREFIX(BIO_write)
 #define BIO_write_all BORINGSSL_ADD_PREFIX(BIO_write_all)
+#define BIO_write_ex BORINGSSL_ADD_PREFIX(BIO_write_ex)
 #define BIO_write_filename BORINGSSL_ADD_PREFIX(BIO_write_filename)
 #define BLAKE2B256 BORINGSSL_ADD_PREFIX(BLAKE2B256)
 #define BLAKE2B256_Final BORINGSSL_ADD_PREFIX(BLAKE2B256_Final)
diff --git a/ssl/bio_ssl.cc b/ssl/bio_ssl.cc
index fc61e9b..7e9d4d6 100644
--- a/ssl/bio_ssl.cc
+++ b/ssl/bio_ssl.cc
@@ -184,8 +184,11 @@
 }
 
 static const BIO_METHOD ssl_method = {
-    BIO_TYPE_SSL, "SSL",   ssl_write, ssl_read,          /*bgets=*/nullptr,
-    ssl_ctrl,     ssl_new, ssl_free,  ssl_callback_ctrl,
+    BIO_TYPE_SSL, "SSL",
+    ssl_write,    /*bwrite_ex=*/nullptr,
+    ssl_read,     /*bgets=*/nullptr,
+    ssl_ctrl,     ssl_new,
+    ssl_free,     ssl_callback_ctrl,
 };
 
 const BIO_METHOD *BIO_f_ssl() { return &ssl_method; }