return (__force u16)cpu_to_be16(thdr->length);
 }
 
+static int csk_mem_free(struct chtls_dev *cdev, struct sock *sk)
+{
+       return (cdev->max_host_sndbuf - sk->sk_wmem_queued);
+}
+
+static int csk_wait_memory(struct chtls_dev *cdev,
+                          struct sock *sk, long *timeo_p)
+{
+       DEFINE_WAIT_FUNC(wait, woken_wake_function);
+       int sndbuf, err = 0;
+       long current_timeo;
+       long vm_wait = 0;
+       bool noblock;
+
+       current_timeo = *timeo_p;
+       noblock = (*timeo_p ? false : true);
+       sndbuf = cdev->max_host_sndbuf;
+       if (csk_mem_free(cdev, sk)) {
+               current_timeo = (prandom_u32() % (HZ / 5)) + 2;
+               vm_wait = (prandom_u32() % (HZ / 5)) + 2;
+       }
+
+       add_wait_queue(sk_sleep(sk), &wait);
+       while (1) {
+               sk_set_bit(SOCKWQ_ASYNC_NOSPACE, sk);
+
+               if (sk->sk_err || (sk->sk_shutdown & SEND_SHUTDOWN))
+                       goto do_error;
+               if (!*timeo_p) {
+                       if (noblock)
+                               set_bit(SOCK_NOSPACE, &sk->sk_socket->flags);
+                       goto do_nonblock;
+               }
+               if (signal_pending(current))
+                       goto do_interrupted;
+               sk_clear_bit(SOCKWQ_ASYNC_NOSPACE, sk);
+               if (csk_mem_free(cdev, sk) && !vm_wait)
+                       break;
+
+               set_bit(SOCK_NOSPACE, &sk->sk_socket->flags);
+               sk->sk_write_pending++;
+               sk_wait_event(sk, ¤t_timeo, sk->sk_err ||
+                             (sk->sk_shutdown & SEND_SHUTDOWN) ||
+                             (csk_mem_free(cdev, sk) && !vm_wait), &wait);
+               sk->sk_write_pending--;
+
+               if (vm_wait) {
+                       vm_wait -= current_timeo;
+                       current_timeo = *timeo_p;
+                       if (current_timeo != MAX_SCHEDULE_TIMEOUT) {
+                               current_timeo -= vm_wait;
+                               if (current_timeo < 0)
+                                       current_timeo = 0;
+                       }
+                       vm_wait = 0;
+               }
+               *timeo_p = current_timeo;
+       }
+do_rm_wq:
+       remove_wait_queue(sk_sleep(sk), &wait);
+       return err;
+do_error:
+       err = -EPIPE;
+       goto do_rm_wq;
+do_nonblock:
+       err = -EAGAIN;
+       goto do_rm_wq;
+do_interrupted:
+       err = sock_intr_errno(*timeo_p);
+       goto do_rm_wq;
+}
+
 int chtls_sendmsg(struct sock *sk, struct msghdr *msg, size_t size)
 {
        struct chtls_sock *csk = rcu_dereference_sk_user_data(sk);
                        copy = mss - skb->len;
                        skb->ip_summed = CHECKSUM_UNNECESSARY;
                }
+               if (!csk_mem_free(cdev, sk))
+                       goto wait_for_sndbuf;
 
                if (is_tls_tx(csk) && !csk->tlshws.txleft) {
                        struct tls_hdr hdr;
                if (ULP_SKB_CB(skb)->flags & ULPCB_FLAG_NO_APPEND)
                        push_frames_if_head(sk);
                continue;
+wait_for_sndbuf:
+               set_bit(SOCK_NOSPACE, &sk->sk_socket->flags);
 wait_for_memory:
-               err = sk_stream_wait_memory(sk, &timeo);
+               err = csk_wait_memory(cdev, sk, &timeo);
                if (err)
                        goto do_error;
        }
                   int offset, size_t size, int flags)
 {
        struct chtls_sock *csk;
+       struct chtls_dev *cdev;
        int mss, err, copied;
        struct tcp_sock *tp;
        long timeo;
        tp = tcp_sk(sk);
        copied = 0;
        csk = rcu_dereference_sk_user_data(sk);
+       cdev = csk->cdev;
        timeo = sock_sndtimeo(sk, flags & MSG_DONTWAIT);
 
        err = sk_stream_wait_connect(sk, &timeo);
                if (!skb || (ULP_SKB_CB(skb)->flags & ULPCB_FLAG_NO_APPEND) ||
                    copy <= 0) {
 new_buf:
+                       if (!csk_mem_free(cdev, sk))
+                               goto wait_for_sndbuf;
 
                        if (is_tls_tx(csk)) {
                                skb = get_record_skb(sk,
                                skb = get_tx_skb(sk, 0);
                        }
                        if (!skb)
-                               goto do_error;
+                               goto wait_for_memory;
                        copy = mss;
                }
                if (copy > size)
                if (unlikely(ULP_SKB_CB(skb)->flags & ULPCB_FLAG_NO_APPEND))
                        push_frames_if_head(sk);
                continue;
-
+wait_for_sndbuf:
                set_bit(SOCK_NOSPACE, &sk->sk_socket->flags);
+wait_for_memory:
+               err = csk_wait_memory(cdev, sk, &timeo);
+               if (err)
+                       goto do_error;
        }
 out:
        csk_reset_flag(csk, CSK_TX_MORE_DATA);