~nickbp/tokio-scgi

ref: f78dff6f33dc8869bff7d23ff9b2c8734b1d2274 tokio-scgi/src/client.rs -rw-r--r-- 3.4 KiB
f78dff6fNick Parker Fix markdown for project link, whoops 2 years ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
#![deny(warnings, rust_2018_idioms)]

use bytes::{BufMut, BytesMut};
use std::io;
use tokio_codec::{Decoder, Encoder};

const NUL: u8 = b'\0';

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SCGIRequest {
    /// The Vec contains the headers. The BytesMut optionally contains raw byte data to include in
    /// the request body.
    Request(Vec<(String, String)>, BytesMut),

    /// Additional body fragment(s) to be used for streaming request data.
    BodyFragment(BytesMut),
}

/// A `Codec` implementation that creates and parses SCGI requests.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SCGICodec {}

impl SCGICodec {
    /// Returns a `SCGIClientCodec` for creating SCGI-format requests.
    pub fn new() -> SCGICodec {
        SCGICodec {}
    }
}

impl Decoder for SCGICodec {
    type Item = BytesMut;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<BytesMut>, io::Error> {
        // Forward content (HTTP response, typically?) as-is
        Ok(Some(buf.split_to(buf.len())))
    }
}

/// Creates and produces SCGI requests. Invoke once with `Request`, followed by zero or more calls
/// with `BodyFragment`.
impl Encoder for SCGICodec {
    type Item = SCGIRequest;
    type Error = io::Error;

    fn encode(&mut self, data: SCGIRequest, buf: &mut BytesMut) -> Result<(), io::Error> {
        match data {
            SCGIRequest::Request(env_map, body) => {
                // Calculate size needed for header netstring
                let mut sum_header_size: usize = 0;
                for (k, v) in &env_map {
                    // While we're iterating over the keys/values, do some basic validation per the
                    // SCGI protocol spec.
                    if k.len() == 0 {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidInput,
                            format!("Keys in request header cannot be empty"),
                        ));
                    }
                    if k.as_bytes().contains(&NUL) || v.as_bytes().contains(&NUL) {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidInput,
                            format!("Keys/values in request header cannot contain NUL character"),
                        ));
                    }
                    // Include 2 x NUL in size:
                    sum_header_size += k.len() + 1/*NUL*/ + v.len() + 1/*NUL*/;
                }
                let netstring_size_str = sum_header_size.to_string();
                // Include ':' and ',' in buffer, not included in netstring size:
                buf.reserve(
                    netstring_size_str.len() + 1/*:*/ + sum_header_size + 1/*,*/ + body.len(),
                );

                // Insert the header content into the reserved buffer.
                buf.put(netstring_size_str);
                buf.put(b':');
                for (k, v) in &env_map {
                    buf.put(k);
                    buf.put(NUL);
                    buf.put(v);
                    buf.put(NUL);
                }
                buf.put(b',');

                // Add any body content after the header
                buf.put(body);
            }
            SCGIRequest::BodyFragment(fragment) => {
                // Forward content as-is
                buf.reserve(fragment.len());
                buf.put(fragment);
            }
        }
        Ok(())
    }
}