aboutsummaryrefslogtreecommitdiff
path: root/src/query.rs
blob: 7cf505042af218c42f6b2384e87766a48b99b634 (plain) (blame)
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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
use std::os::unix::prelude::RawFd;

use nix::sys::socket::{self, AddressFamily, MsgFlags, SockFlag, SockProtocol, SockType};

use crate::{
    error::QueryError,
    nlmsg::{
        nft_nlmsg_maxsize, pad_netlink_object_with_variable_size, NfNetlinkAttribute,
        NfNetlinkObject, NfNetlinkWriter,
    },
    parser::{parse_nlmsg, NlMsg},
    sys::{NLM_F_DUMP, NLM_F_MULTI},
    ProtocolFamily,
};

pub(crate) fn recv_and_process<'a, T>(
    sock: RawFd,
    max_seq: Option<u32>,
    cb: Option<&dyn Fn(&[u8], &mut T) -> Result<(), QueryError>>,
    working_data: &'a mut T,
) -> Result<(), QueryError> {
    let mut msg_buffer = vec![0; 2 * nft_nlmsg_maxsize() as usize];
    let mut buf_start = 0;
    let mut end_pos = 0;

    loop {
        let nb_recv = socket::recv(sock, &mut msg_buffer[end_pos..], MsgFlags::empty())
            .map_err(QueryError::NetlinkRecvError)?;
        if nb_recv <= 0 {
            return Ok(());
        }
        end_pos += nb_recv;
        loop {
            let buf = &msg_buffer.as_slice()[buf_start..end_pos];
            // exit the loop and try to receive further messages when we consumed all the buffer
            if buf.len() == 0 {
                break;
            }

            debug!("Calling parse_nlmsg");
            let (nlmsghdr, msg) = parse_nlmsg(&buf)?;
            debug!("Got a valid netlink message: {:?} {:?}", nlmsghdr, msg);

            match msg {
                NlMsg::Done => {
                    return Ok(());
                }
                NlMsg::Error(e) => {
                    if e.error != 0 {
                        return Err(QueryError::NetlinkError(e));
                    }
                }
                NlMsg::Noop => {}
                NlMsg::NfGenMsg(_genmsg, _data) => {
                    if let Some(cb) = cb {
                        cb(&buf[0..nlmsghdr.nlmsg_len as usize], working_data)?;
                    }
                }
            }

            // we cannot know when a sequence of messages will end if the messages do not end
            // with an NlMsg::Done marker while if a maximum sequence number wasn't specified
            if max_seq.is_none() && nlmsghdr.nlmsg_flags & NLM_F_MULTI as u16 == 0 {
                return Err(QueryError::UndecidableMessageTermination);
            }

            // retrieve the next message
            if let Some(max_seq) = max_seq {
                if nlmsghdr.nlmsg_seq >= max_seq {
                    return Ok(());
                }
            }

            // netlink messages are 4bytes aligned
            let aligned_length = pad_netlink_object_with_variable_size(nlmsghdr.nlmsg_len as usize);
            buf_start += aligned_length;
        }
        // Ensure that we always have nft_nlmsg_maxsize() free space available in the buffer.
        // We achieve this by relocating the buffer content at the beginning of the buffer
        if end_pos >= nft_nlmsg_maxsize() as usize {
            if buf_start < end_pos {
                unsafe {
                    std::ptr::copy(
                        msg_buffer[buf_start..end_pos].as_ptr(),
                        msg_buffer.as_mut_ptr(),
                        end_pos - buf_start,
                    );
                }
            }
            end_pos = end_pos - buf_start;
            buf_start = 0;
        }
    }
}

pub(crate) fn socket_close_wrapper<E>(
    sock: RawFd,
    cb: impl FnOnce(RawFd) -> Result<(), E>,
) -> Result<(), QueryError>
where
    QueryError: From<E>,
{
    let ret = cb(sock);

    // we don't need to shutdown the socket (in fact, Linux doesn't support that operation;
    // and return EOPNOTSUPP if we try)
    nix::unistd::close(sock).map_err(QueryError::CloseFailed)?;

    Ok(ret?)
}

/// Returns a buffer containing a netlink message which requests a list of all the netfilter
/// matching objects (e.g. tables, chains, rules, ...).
/// Supply the type of objects to retrieve (e.g. libc::NFT_MSG_GETTABLE), and a search filter.
pub fn get_list_of_objects<T: NfNetlinkAttribute>(
    msg_type: u16,
    seq: u32,
    filter: Option<&T>,
) -> Result<Vec<u8>, QueryError> {
    let mut buffer = Vec::new();
    let mut writer = NfNetlinkWriter::new(&mut buffer);
    writer.write_header(
        msg_type,
        ProtocolFamily::Unspec,
        NLM_F_DUMP as u16,
        seq,
        None,
    );
    if let Some(filter) = filter {
        let buf = writer.add_data_zeroed(filter.get_size());
        unsafe {
            filter.write_payload(buf.as_mut_ptr());
        }
    }
    writer.finalize_writing_object();
    Ok(buffer)
}

/// Lists objects of a certain type (e.g. libc::NFT_MSG_GETTABLE) with the help of a helper
/// function called by mnl::cb_run2.
/// The callback expects a tuple of additional data (supplied as an argument to this function)
/// and of the output vector, to which it should append the parsed object it received.
pub fn list_objects_with_data<'a, Object, Accumulator>(
    data_type: u16,
    cb: &dyn Fn(Object, &mut Accumulator) -> Result<(), QueryError>,
    filter: Option<&Object>,
    working_data: &'a mut Accumulator,
) -> Result<(), QueryError>
where
    Object: NfNetlinkObject + NfNetlinkAttribute,
{
    debug!("Listing objects of kind {}", data_type);
    let sock = socket::socket(
        AddressFamily::Netlink,
        SockType::Raw,
        SockFlag::empty(),
        SockProtocol::NetlinkNetFilter,
    )
    .map_err(QueryError::NetlinkOpenError)?;

    let seq = 0;

    let chains_buf = get_list_of_objects(data_type, seq, filter)?;
    socket::send(sock, &chains_buf, MsgFlags::empty()).map_err(QueryError::NetlinkSendError)?;

    socket_close_wrapper(sock, move |sock| {
        // the kernel should return NLM_F_MULTI objects
        recv_and_process(
            sock,
            None,
            Some(&|buf: &[u8], working_data: &mut Accumulator| {
                debug!("Calling Object::deserialize()");
                cb(Object::deserialize(buf)?.0, working_data)
            }),
            working_data,
        )
    })
}