brpc/include/butil/details/extended_endpoint.hpp

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#ifndef BUTIL_DETAILS_EXTENDED_ENDPOINT_H
#define BUTIL_DETAILS_EXTENDED_ENDPOINT_H

#include <arpa/inet.h>
#include <netdb.h>
#include <unistd.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/un.h>
#include <mutex>
#include <unordered_set>
#include "butil/endpoint.h"
#include "butil/logging.h"
#include "butil/strings/string_piece.h"
#include "butil/resource_pool.h"
#include "butil/memory/singleton_on_pthread_once.h"

namespace butil {
namespace details {

#if __cplusplus >= 201103L
static_assert(sizeof(EndPoint) == sizeof(EndPoint::ip) + sizeof(EndPoint::port),
        "EndPoint size mismatch with the one in POD-style, may cause ABI problem");
#endif

// For ipv6/unix socket address.
//
// We have to keep butil::EndPoint ABI compatible because it is used so widely, and the size of butil::EndPoint is
// too small to store more information such as ipv6 address.
// We store enough information about endpoint in such tiny struct by putting real things in another big object
// holding by ResourcePool. The EndPoint::ip saves ResourceId, while EndPoint::port denotes if the EndPoint object
// is an old style ipv4 endpoint.
// Note that since ResourcePool has been implemented in bthread, we copy it into this repo and change its namespace to
// butil::details. Those two headers will not be published.

// If EndPoint.port equals to this value, we should get the extended endpoint in resource pool.
const static int EXTENDED_ENDPOINT_PORT = 123456789;

class ExtendedEndPoint;

// A global unordered set to dedup ExtendedEndPoint
// ExtendedEndPoints which have same ipv6/unix socket address must have same id,
// so that user can simply use the value of EndPoint for comparision.
class GlobalEndPointSet {
public:
    ExtendedEndPoint* insert(ExtendedEndPoint* p);

    void erase(ExtendedEndPoint* p);

    static GlobalEndPointSet* instance() {
        return ::butil::get_leaky_singleton<GlobalEndPointSet>();
    }

private:
    struct Hash {
        size_t operator()(ExtendedEndPoint* const& p) const;
    };

    struct Equals {
        bool operator()(ExtendedEndPoint* const& p1, ExtendedEndPoint* const& p2) const;
    };

    typedef std::unordered_set<ExtendedEndPoint*, Hash, Equals> SetType;
    SetType _set;
    std::mutex _mutex;
};

class ExtendedEndPoint {
public:
    // Construct ExtendedEndPoint.
    // User should use create() functions to get ExtendedEndPoint instance.
    ExtendedEndPoint(void) {
        _ref_count.store(0, butil::memory_order_relaxed);
        _u.sa.sa_family = AF_UNSPEC;
    }

public:
    // Create ExtendedEndPoint.
    // If creation is successful, create()s will embed the ExtendedEndPoint instance in the given EndPoint*,
    // and return it as well. Or else, the given EndPoint* won't be touched.
    //
    // The format of the parameter is inspired by nginx.
    // Valid forms are:
    // - ipv6
    // without port: [2400:da00::3b0b]
    // with    port: [2400:da00::3b0b]:8080
    // - unix domain socket
    // abslute path : unix:/path/to/file.sock
    // relative path: unix:path/to/file.sock

    static ExtendedEndPoint* create(StringPiece sp, EndPoint* ep) {
        sp.trim_spaces();
        if (sp.empty()) {
            return NULL;
        }
        if (sp[0] == '[') {
            size_t colon_pos = sp.find(']');
            if (colon_pos == StringPiece::npos || colon_pos == 1 /* [] is invalid */ || ++colon_pos >= sp.size()) {
                return NULL;
            }
            StringPiece port_sp = sp.substr(colon_pos);
            if (port_sp.size() < 2 /* colon and at least one integer */ || port_sp[0] != ':') {
                return NULL;
            }
            port_sp.remove_prefix(1); // remove `:'
            if (port_sp.size() > 5) { // max 65535
                return NULL;
            }
            char buf[6];
            buf[port_sp.copy(buf, port_sp.size())] = '\0';
            char* end = NULL;
            int port = ::strtol(buf, &end, 10 /* base */);
            if (end != buf + port_sp.size()) {
                return NULL;
            }
            return create(sp.substr(0, colon_pos), port, ep);
        } else if (sp.starts_with("unix:")) {
            return create(sp, EXTENDED_ENDPOINT_PORT, ep);
        }
        return NULL;
    }

    static ExtendedEndPoint* create(StringPiece sp, int port, EndPoint* ep) {
        sp.trim_spaces();
        if (sp.empty()) {
            return NULL;
        }
        ExtendedEndPoint* eep = NULL;
        if (sp[0] == '[' && port >= 0 && port <= 65535) {
            if (sp.back() != ']' || sp.size() == 2 || sp.size() - 2 >= INET6_ADDRSTRLEN) {
                return NULL;
            }
            char buf[INET6_ADDRSTRLEN];
            buf[sp.copy(buf, sp.size() - 2 /* skip `[' and `]' */, 1 /* skip `[' */)] = '\0';

            in6_addr addr;
            if (inet_pton(AF_INET6, buf, &addr) != 1 /* succ */) {
                return NULL;
            }

            eep = new_extended_endpoint(AF_INET6);
            if (eep) {
                eep->_u.in6.sin6_addr = addr;
                eep->_u.in6.sin6_port = htons(port);
                eep->_u.in6.sin6_flowinfo = 0u;
                eep->_u.in6.sin6_scope_id = 0u;
                eep->_socklen = sizeof(_u.in6);
#if defined(OS_MACOSX)
                eep->_u.in6.sin6_len = eep->_socklen;
#endif
            }
        } else if (sp.starts_with("unix:")) { // ignore port
            sp.remove_prefix(5); // remove `unix:'
            if (sp.empty() || sp.size() >= UDS_PATH_SIZE) {
                return NULL;
            }
            eep = new_extended_endpoint(AF_UNIX);
            if (eep) {
                int size = sp.copy(eep->_u.un.sun_path, sp.size());
                eep->_u.un.sun_path[size] = '\0';
                eep->_socklen = offsetof(sockaddr_un, sun_path) + size + 1;
#if defined(OS_MACOSX)
                eep->_u.un.sun_len = eep->_socklen;
#endif
            }
        }
        if (eep) {
            eep = dedup(eep);
            eep->embed_to(ep);
        }
        return eep;
    }

    static ExtendedEndPoint* create(sockaddr_storage* ss, socklen_t size, EndPoint* ep) {
        ExtendedEndPoint* eep = NULL;
        if (ss->ss_family == AF_INET6 || ss->ss_family == AF_UNIX) {
            eep = new_extended_endpoint(ss->ss_family);
        }
        if (eep) {
            memcpy(&eep->_u.ss, ss, size);
            eep->_socklen = size;
            if (ss->ss_family == AF_UNIX && size == offsetof(sockaddr_un, sun_path)) {
                // See unix(7): When the address of an unnamed socket is returned, 
                // its length is sizeof(sa_family_t), and sun_path should not be inspected.
                eep->_u.un.sun_path[0] = '\0';
            }
            eep = dedup(eep);
            eep->embed_to(ep);
        }
        return eep;
    }

    // Get ExtendedEndPoint instance from EndPoint
    static ExtendedEndPoint* address(const EndPoint& ep) {
        if (!is_extended(ep)) {
            return NULL;
        }
        ::butil::ResourceId<ExtendedEndPoint> id;
        id.value = ep.ip.s_addr;
        ExtendedEndPoint* eep = ::butil::address_resource<ExtendedEndPoint>(id);
        CHECK(eep) << "fail to address ExtendedEndPoint from EndPoint";
        return eep;
    }

    // Check if an EndPoint has embedded ExtendedEndPoint
    static bool is_extended(const butil::EndPoint& ep) {
        return ep.port == EXTENDED_ENDPOINT_PORT;
    }

private:
    friend class GlobalEndPointSet;

    static GlobalEndPointSet* global_set() {
        return GlobalEndPointSet::instance();
    }

    static ExtendedEndPoint* new_extended_endpoint(sa_family_t family) {
        ::butil::ResourceId<ExtendedEndPoint> id;
        ExtendedEndPoint* eep = ::butil::get_resource(&id);
        if (eep) {
            int64_t old_ref = eep->_ref_count.load(butil::memory_order_relaxed);
            CHECK(old_ref == 0) << "new ExtendedEndPoint has reference " << old_ref;
            CHECK(eep->_u.sa.sa_family == AF_UNSPEC) << "new ExtendedEndPoint has family " << eep->_u.sa.sa_family << " set";
            eep->_ref_count.store(1, butil::memory_order_relaxed);
            eep->_id = id;
            eep->_u.sa.sa_family = family;
        }
        return eep;
    }

    void embed_to(EndPoint* ep) const {
        CHECK(0 == _id.value >> 32) << "ResourceId beyond index";
        ep->reset();
        ep->ip = ip_t{static_cast<uint32_t>(_id.value)};
        ep->port = EXTENDED_ENDPOINT_PORT;
    }

    static ExtendedEndPoint* dedup(ExtendedEndPoint* eep) {
        eep->_hash = std::hash<std::string>()(std::string((const char*)&eep->_u, eep->_socklen));

        ExtendedEndPoint* first_eep = global_set()->insert(eep);
        if (first_eep != eep) {
            eep->_ref_count.store(0, butil::memory_order_relaxed);
            eep->_u.sa.sa_family = AF_UNSPEC;
            ::butil::return_resource(eep->_id);
        }
        return first_eep;
    }

public:

    void dec_ref(void) {
        int64_t old_ref = _ref_count.fetch_sub(1, butil::memory_order_relaxed);
        CHECK(old_ref >= 1) << "ExtendedEndPoint has unexpected reference " << old_ref;
        if (old_ref == 1) {
            global_set()->erase(this);
            _u.sa.sa_family = AF_UNSPEC;
            ::butil::return_resource(_id);
        }
    }

    void inc_ref(void) {
        int64_t old_ref = _ref_count.fetch_add(1, butil::memory_order_relaxed);
        CHECK(old_ref >= 1) << "ExtendedEndPoint has unexpected reference " << old_ref;
    }

    sa_family_t family(void) const {
        return _u.sa.sa_family;
    }

    int to(sockaddr_storage* ss) const {
        memcpy(ss, &_u.ss, _socklen);
        return _socklen;
    }

    void to(EndPointStr* ep_str) const {
        if (_u.sa.sa_family == AF_UNIX) {
            snprintf(ep_str->_buf, sizeof(ep_str->_buf), "unix:%s", _u.un.sun_path);
        } else if (_u.sa.sa_family == AF_INET6) {
            char buf[INET6_ADDRSTRLEN] = {0};
            const char* ret = inet_ntop(_u.sa.sa_family, &_u.in6.sin6_addr, buf, sizeof(buf));
            CHECK(ret) << "fail to do inet_ntop";
            snprintf(ep_str->_buf, sizeof(ep_str->_buf), "[%s]:%d", buf, ntohs(_u.in6.sin6_port));
        } else {
            CHECK(0) << "family " << _u.sa.sa_family << " not supported";
        }
    }

    int to_hostname(char* host, size_t host_len) const {
        if (_u.sa.sa_family == AF_UNIX) {
            snprintf(host, host_len, "unix:%s", _u.un.sun_path);
            return 0;
        } else if (_u.sa.sa_family == AF_INET6) {
            sockaddr_in6 sa = _u.in6;
            if (getnameinfo((const sockaddr*) &sa, sizeof(sa), host, host_len, NULL, 0, NI_NAMEREQD) != 0) {
                return -1;
            }
            size_t len = ::strlen(host);
            if (len + 1 < host_len) {
                snprintf(host + len, host_len - len, ":%d", _u.in6.sin6_port);
            }
            return 0;
        } else {
            CHECK(0) << "family " << _u.sa.sa_family << " not supported";
            return -1;
        }
    }

private:
    static const size_t UDS_PATH_SIZE = sizeof(sockaddr_un::sun_path);

    butil::atomic<int64_t> _ref_count;
    butil::ResourceId<ExtendedEndPoint> _id;
    size_t _hash;  // pre-compute hash code of sockaddr for saving unordered_set query time
    socklen_t _socklen; // valid data length of sockaddr
    union {
        sockaddr sa;
        sockaddr_in6 in6;
        sockaddr_un un;
        sockaddr_storage ss;
    } _u;
};

inline ExtendedEndPoint* GlobalEndPointSet::insert(ExtendedEndPoint* p) {
    std::unique_lock<std::mutex> lock(_mutex);
    auto it = _set.find(p);
    if (it != _set.end()) {
        if ((*it)->_ref_count.fetch_add(1, butil::memory_order_relaxed) == 0) {
            // another thread is calling dec_ref(), do not reuse it
            (*it)->_ref_count.fetch_sub(1, butil::memory_order_relaxed);
            _set.erase(it);
            _set.insert(p);
            return p;
        } else {
            // the ExtendedEndPoint is valid, reuse it 
            return *it;
        }
    }
    _set.insert(p);
    return p;
}

inline void GlobalEndPointSet::erase(ExtendedEndPoint* p) {
    std::unique_lock<std::mutex> lock(_mutex);
    auto it = _set.find(p);
    if (it == _set.end() || *it != p) {
        // another thread has been erase it
        return;
    }
    _set.erase(it);
}

inline size_t GlobalEndPointSet::Hash::operator()(ExtendedEndPoint* const& p) const {
    return p->_hash;
}

inline bool GlobalEndPointSet::Equals::operator()(ExtendedEndPoint* const& p1, ExtendedEndPoint* const& p2) const {
    return p1->_socklen == p2->_socklen 
            && memcmp(&p1->_u, &p2->_u, p1->_socklen) == 0;
}

} // namespace details
} // namespace butil

#endif // BUTIL_DETAILS_EXTENDED_ENDPOINT_H