Skip to content

Instantly share code, notes, and snippets.

@vans163

vans163/sock_nif

Created Dec 18, 2018
Embed
What would you like to do?
//gcc -O3 -fpic -shared -I/usr/lib/erlang/usr/include/ nif_tcp.c -o nif_tcp.so
#include "erl_nif.h"
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <stdio.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <stdlib.h>
#include <fcntl.h> /* Added for the nonblocking socket */
int accept4(int sockfd, struct sockaddr *addr, socklen_t *addrlen, int flags);
ERL_NIF_TERM mk_atom(ErlNifEnv* env, const char* atom)
{
ERL_NIF_TERM ret;
if(!enif_make_existing_atom(env, atom, &ret, ERL_NIF_LATIN1))
return enif_make_atom(env, atom);
return ret;
}
ERL_NIF_TERM mk_error(ErlNifEnv* env, const char* mesg)
{
return enif_make_tuple2(env, mk_atom(env, "error"), mk_atom(env, mesg));
}
static ERL_NIF_TERM nif_listen(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
struct sockaddr_un address;
address.sun_family = AF_UNIX;
strcpy(address.sun_path, "/tmp/star.sock");
unlink("/tmp/star.sock");
int fd = socket(AF_UNIX, SOCK_STREAM, 0);
bind(fd, (struct sockaddr*)(&address), sizeof(address));
listen(fd, 1000);
//int i = 1;
//setsockopt( fd, IPPROTO_TCP, TCP_NODELAY, (void *)&i, sizeof(i));
//setsockopt( fd, IPPROTO_TCP, TCP_QUICKACK, (void *)&i, sizeof(i));
//non-blocking accept
fcntl(fd, F_SETFL, O_NONBLOCK);
return enif_make_tuple2(env, mk_atom(env, "ok"), enif_make_int(env, fd));
}
static ERL_NIF_TERM nif_accept(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
int sockfd;
if (!enif_get_int(env, argv[0], &sockfd))
return mk_error(env, "not_a_number");
int client_fd = accept4(sockfd, NULL, NULL, SOCK_NONBLOCK);
if (client_fd == -1) {
if (errno == EWOULDBLOCK) {
return enif_make_tuple2(env, mk_atom(env, "error"), mk_atom(env, "ewouldblock"));
}
if (errno == EAGAIN) {
return enif_make_tuple2(env, mk_atom(env, "error"), mk_atom(env, "eagain"));
}
return enif_make_tuple2(env, mk_atom(env, "error"), enif_make_int(env, errno));
} else {
return enif_make_tuple2(env, mk_atom(env, "ok"), enif_make_int(env, client_fd));
}
}
static ERL_NIF_TERM nif_send(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
int sockfd;
if (!enif_get_int(env, argv[0], &sockfd))
return mk_error(env, "not_a_number");
ErlNifBinary blob;
if (!enif_inspect_binary(env, argv[1], &blob))
return mk_error(env, "not_a_blob_binary");
int res = send(sockfd, blob.data, blob.size, 0);
if (res == -1) {
if (errno == EWOULDBLOCK) {
return enif_make_tuple2(env, mk_atom(env, "error"), mk_atom(env, "ewouldblock"));
}
if (errno == EAGAIN) {
return enif_make_tuple2(env, mk_atom(env, "error"), mk_atom(env, "eagain"));
}
return enif_make_tuple2(env, mk_atom(env, "error"), enif_make_int(env, errno));
} else if (res != blob.size) {
return enif_make_tuple2(env, mk_atom(env, "partial_send"), enif_make_int(env, blob.size - res));
} else {
return mk_atom(env, "ok");
}
}
static ERL_NIF_TERM nif_recv(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
int sockfd;
if (!enif_get_int(env, argv[0], &sockfd))
return mk_error(env, "not_a_number");
char buf[1024];
int res = recv(sockfd, buf, 1024, 0);
if (res < 1) {
if (errno == EWOULDBLOCK) {
return enif_make_tuple2(env, mk_atom(env, "error"), mk_atom(env, "ewouldblock"));
}
if (errno == EAGAIN) {
return enif_make_tuple2(env, mk_atom(env, "error"), mk_atom(env, "eagain"));
}
return enif_make_tuple2(env, mk_atom(env, "error"), enif_make_int(env, errno));
} else {
ErlNifBinary bin_output;
enif_alloc_binary(res, &bin_output);
memcpy(bin_output.data, buf, res);
return enif_make_binary(env, &bin_output);
}
}
int upgrade(ErlNifEnv* env, void** priv_data, void** old_priv_data, ERL_NIF_TERM load_info)
{
return 0;
}
static ErlNifFunc nif_funcs[] = {
{"listen", 0, nif_listen},
{"accept", 1, nif_accept},
{"send", 2, nif_send},
{"recv", 1, nif_recv},
};
ERL_NIF_INIT(Elixir.Stargate.NifTcp, nif_funcs, NULL, NULL, &upgrade, NULL)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment