<commit_before>// Copyright (c) 2009-2010 Satoshi Nakamoto
// Copyright (c) 2009-2012 The Bitcoin developers
// Distributed under the MIT/X11 software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "irc.h"
#include "net.h"
#include "base58.h"
#include "init.h"
#include "chainparams.h"

#include <stdint.h>
#define __STDC_FORMAT_MACROS
#include <inttypes.h>

size_t strlcpy(char *dst, const char *src, size_t size) {
	size_t srclen;
    size--;
	srclen = strlen(src);
	if (srclen > size) {
    	srclen = size;
	}
	memcpy(dst, src, srclen);
  	dst[srclen] = '\0';
	return (srclen);
}

using namespace std;
using namespace boost;

int nGotIRCAddresses = 0;

#pragma pack(push, 1)
struct ircaddr {
    struct in_addr ip;
    short port;
};
#pragma pack(pop)

bool RecvIRCLine(SOCKET hSocket, string& strLine) {
    strLine = "";
    for (;;) {
    	boost::this_thread::interruption_point();
    	if (ShutdownRequested())
            return false;
        char c;
        int nBytes = recv(hSocket, &c, 1, 0);
        if (nBytes > 0) {
            if (c == '\n')
                continue;
            if (c == '\r')
                return true;
            strLine += c;
            if (strLine.size() >= 9000)
                return true;
        } else if (nBytes <= 0) {
            if (nBytes < 0) {
                int nErr = WSAGetLastError();
                if (nErr == WSAEMSGSIZE)
                    continue;
                if (nErr == WSAEWOULDBLOCK || nErr == WSAEINTR || nErr == WSAEINPROGRESS) {
                    MilliSleep(40);
                    continue;
                }
            }
            if (!strLine.empty())
                return true;
            if (nBytes == 0) {
                // socket closed
                LogPrintf("IRC: socket closed\n");
                return false;
            } else {
                // socket error
                int nErr = WSAGetLastError();
                LogPrintf("IRC: recv failed: %d\n", nErr);
                return false;
            }
        }
    }
}

string EncodeAddress(const CService& addr) {
    struct ircaddr tmp;
    if (addr.GetInAddr(&tmp.ip))
    {
        tmp.port = htons(addr.GetPort());

        vector<unsigned char> vch(UBEGIN(tmp), UEND(tmp));
        return string("u") + EncodeBase58Check(vch);
    }
    return "";
}

bool DecodeAddress(string str, CService& addr) {
    vector<unsigned char> vch;
    if (!DecodeBase58Check(str.substr(1), vch))
        return false;

    struct ircaddr tmp;
    if (vch.size() != sizeof(tmp))
        return false;
    memcpy(&tmp, &vch[0], sizeof(tmp));

    addr = CService(tmp.ip, ntohs(tmp.port));
    return true;
}

static bool Send(SOCKET hSocket, const char* pszSend) {
    if (strstr(pszSend, "PONG") != pszSend) {
        printf("IRC SENDING: %s\n", pszSend);
    }
    const char* psz = pszSend;
    const char* pszEnd = psz + strlen(psz);
    while (psz < pszEnd) {
        int ret = send(hSocket, psz, pszEnd - psz, MSG_NOSIGNAL);
        if (ret < 0)
            return false;
        psz += ret;
    }
    return true;
}

bool RecvLineIRC(SOCKET hSocket, string& strLine) {
    for (;;) {
    	boost::this_thread::interruption_point();
    	if (ShutdownRequested())
            return false;
        bool fRet = RecvIRCLine(hSocket, strLine);
        if (fRet) {
            vector<string> vWords;
            ParseString(strLine, ' ', vWords);
            if (vWords.size() >= 1 && vWords[0] == "PING") {
                strLine[1] = 'O';
                strLine += '\r';
                Send(hSocket, strLine.c_str());
                continue;
            }
        }
        return fRet;
    }
}

int RecvUntil(SOCKET hSocket, const char* psz1, const char* psz2=NULL, const char* psz3=NULL, const char* psz4=NULL) {
    for (;;) {
    	boost::this_thread::interruption_point();
        string strLine;
        strLine.reserve(10000);
        if (!RecvLineIRC(hSocket, strLine))
            return 0;
        LogPrintf("IRC %s\n", strLine.c_str());
        if (psz1 && strLine.find(psz1) != string::npos)
            return 1;
        if (psz2 && strLine.find(psz2) != string::npos)
            return 2;
        if (psz3 && strLine.find(psz3) != string::npos)
            return 3;
        if (psz4 && strLine.find(psz4) != string::npos)
            return 4;
    }
}

bool Wait(int nSeconds) {
    if (ShutdownRequested()) {
        return false;
    }
    LogPrintf("IRC waiting %d seconds to reconnect\n", nSeconds);
    for (int i = 0; i < nSeconds; i++) {
    	boost::this_thread::interruption_point();
        if (ShutdownRequested()) {
            return false;
        }
        MilliSleep(1000);
    }
    return true;
}

bool RecvCodeLine(SOCKET hSocket, const char* psz1, string& strRet) {
    strRet.clear();
    for (;;) {
    	boost::this_thread::interruption_point();
        string strLine;
        if (!RecvLineIRC(hSocket, strLine)) {
            return false;
        }

        vector<string> vWords;
        ParseString(strLine, ' ', vWords);
        if (vWords.size() < 2) {
            continue;
        }

        if (vWords[1] == psz1) {
            LogPrintf("IRC %s\n", strLine.c_str());
            strRet = strLine;
            return true;
        }
    }
}

bool GetIPFromIRC(SOCKET hSocket, string strMyName, CNetAddr& ipRet) {
    Send(hSocket, strprintf("USERHOST %s\r", strMyName.c_str()).c_str());

    string strLine;
    if (!RecvCodeLine(hSocket, "302", strLine)) {
        return false;
    }

    vector<string> vWords;
    ParseString(strLine, ' ', vWords);
    if (vWords.size() < 4) {
        return false;
    }

    string str = vWords[3];
    if (str.rfind("@") == string::npos) {
        return false;
    }
    string strHost = str.substr(str.rfind("@")+1);

    // Hybrid IRC used by lfnet always returns IP when you userhost yourself,
    // but in case another IRC is ever used this should work.
    LogPrintf("GetIPFromIRC() got userhost %s\n", strHost.c_str());
    CNetAddr addr(strHost, true);
    if (!addr.IsValid()) {
        return false;
    }
    ipRet = addr;

    return true;
}

void ThreadIRCSeed() {
    // Make this thread recognisable as the IRC seeding thread
    RenameThread("Ember-ircseed");

    int nErrorWait = 10;
    int nRetryWait = 10;
    int nNameRetry = 0;

begin_irc:
	boost::this_thread::interruption_point();
    CService addrConnect("62.210.131.147", 7777); // irc.lfnet.org

    CService addrIRC("irc.lfnet.org", 7777, true);
    if (addrIRC.IsValid()) {
        addrConnect = addrIRC;
    }

    SOCKET hSocket;
    if (!ConnectSocket(addrConnect, hSocket, 60*5)) {
        LogPrintf("IRC connect failed\n");
        nErrorWait = nErrorWait * 11 / 10;
        if (Wait(nErrorWait += 60))
            goto begin_irc;
        else
            return;
    }

    if (!RecvUntil(hSocket, "Found your hostname", "using your IP address instead", "Couldn't look up your hostname", "ignoring hostname")) {
        closesocket(hSocket);
        hSocket = INVALID_SOCKET;
        nErrorWait = nErrorWait * 11 / 10;
        if (Wait(nErrorWait += 60))
            goto begin_irc;
        else
            return;
    }

    CNetAddr addrIPv4("1.2.3.4"); // arbitrary IPv4 address to make GetLocal prefer IPv4 addresses
    CService addrLocal;
    string strMyName;
    // Don't use our IP as our nick if we're not listening
    // or if it keeps failing because the nick is already in use.
    if (!fNoListen && GetLocal(addrLocal, &addrIPv4) && nNameRetry<3)
        strMyName = EncodeAddress(GetLocalAddress(&addrConnect));
    if (strMyName == "")
        strMyName = strprintf("x%llu", GetRand(1000000000));

    Send(hSocket, strprintf("NICK %s\r", strMyName.c_str()).c_str());
    Send(hSocket, strprintf("USER %s 8 * : %s\r", strMyName.c_str(), strMyName.c_str()).c_str());

    int nRet = RecvUntil(hSocket, " 004 ", " 433 ");
    if (nRet != 1)
    {
        closesocket(hSocket);
        hSocket = INVALID_SOCKET;
        if (nRet == 2) {
            LogPrintf("IRC name already in use\n");
            nNameRetry++;
            Wait(10);
            goto begin_irc;
        }
        nErrorWait = nErrorWait * 11 / 10;
        if (Wait(nErrorWait += 60))
            goto begin_irc;
        else
            return;
    }
    nNameRetry = 0;

    boost::this_thread::interruption_point();
    MilliSleep(500);
    boost::this_thread::interruption_point();

    // Get our external IP from the IRC server and re-nick before joining the channel
    CNetAddr addrFromIRC;
    if (GetIPFromIRC(hSocket, strMyName, addrFromIRC)) {
        LogPrintf("GetIPFromIRC() returned %s\n", addrFromIRC.ToString().c_str());
        // Don't use our IP as our nick if we're not listening
        if (!fNoListen && addrFromIRC.IsRoutable()) {
            // IRC lets you to re-nick
            AddLocal(addrFromIRC, LOCAL_IRC);
            strMyName = EncodeAddress(GetLocalAddress(&addrConnect));
            Send(hSocket, strprintf("NICK %s\r", strMyName.c_str()).c_str());
        }
    }

    if (TestNet()) {
        Send(hSocket, "JOIN #Ember0T\r");
        Send(hSocket, "WHO #Ember0T\r");
    } else {
        // randomly join #Ember00-#Ember05
        // int channel_number = GetRandInt(5);

        // Channel number is always 0 for initial release
        int channel_number = 0;
        Send(hSocket, strprintf("JOIN #Ember%02d\r", channel_number).c_str());
        Send(hSocket, strprintf("WHO #Ember%02d\r", channel_number).c_str());
    }

    int64 nStart = GetTime();
    string strLine;
    strLine.reserve(10000);

    while (RecvLineIRC(hSocket, strLine)) {

    	boost::this_thread::interruption_point();

        if (strLine.empty() || strLine.size() > 900 || strLine[0] != ':')
            continue;

        vector<string> vWords;
        ParseString(strLine, ' ', vWords);
        if (vWords.size() < 2)
            continue;

        char pszName[10000];
        pszName[0] = '\0';

        if (vWords[1] == "352" && vWords.size() >= 8) {
            // index 7 is limited to 16 characters
            // could get full length name at index 10, but would be different from join messages
            strlcpy(pszName, vWords[7].c_str(), sizeof(pszName));
            //LogPrintf("IRC got who\n");
        }

        if (vWords[1] == "JOIN" && vWords[0].size() > 1) {
            // :username!username@50000007.F000000B.90000002.IP JOIN :#channelname
            strlcpy(pszName, vWords[0].c_str() + 1, sizeof(pszName));
            if (strchr(pszName, '!'))
                *strchr(pszName, '!') = '\0';
            //LogPrintf("IRC got join\n");
        }

        if (pszName[0] == 'u') {
            CAddress addr;
            if (DecodeAddress(pszName, addr)) {
                addr.nTime = GetAdjustedTime();
                if (addrman.Add(addr, addrConnect, 51 * 60))
                    LogPrintf("IRC got new address: %s\n", addr.ToString().c_str());
                nGotIRCAddresses++;
            } else {
                LogPrintf("IRC decode failed\n");
            }
        }
    }
    closesocket(hSocket);
    hSocket = INVALID_SOCKET;
}

#ifdef TEST
int main(int argc, char *argv[])
{
    WSADATA wsadata;
    if (WSAStartup(MAKEWORD(2,2), &wsadata) != NO_ERROR) {
        LogPrintf("Error at WSAStartup()\n");
        return false;
    }

    ThreadIRCSeed(NULL);

    WSACleanup();
    return 0;
}
#endif<commit_msg>Change Sleep to sleep_for crossplatform alternative<commit_after>// Copyright (c) 2009-2010 Satoshi Nakamoto
// Copyright (c) 2009-2012 The Bitcoin developers
// Distributed under the MIT/X11 software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "irc.h"
#include "net.h"
#include "base58.h"
#include "init.h"
#include "chainparams.h"

#include <stdint.h>
#define __STDC_FORMAT_MACROS
#include <inttypes.h>
#include <chrono>

size_t strlcpy(char *dst, const char *src, size_t size) {
	size_t srclen;
    size--;
	srclen = strlen(src);
	if (srclen > size) {
    	srclen = size;
	}
	memcpy(dst, src, srclen);
  	dst[srclen] = '\0';
	return (srclen);
}

using namespace std;
using namespace boost;

int nGotIRCAddresses = 0;

#pragma pack(push, 1)
struct ircaddr {
    struct in_addr ip;
    short port;
};
#pragma pack(pop)

bool RecvIRCLine(SOCKET hSocket, string& strLine) {
    strLine = "";
    for (;;) {
    	boost::this_thread::interruption_point();
    	if (ShutdownRequested())
            return false;
        char c;
        int nBytes = recv(hSocket, &c, 1, 0);
        if (nBytes > 0) {
            if (c == '\n')
                continue;
            if (c == '\r')
                return true;
            strLine += c;
            if (strLine.size() >= 9000)
                return true;
        } else if (nBytes <= 0) {
            if (nBytes < 0) {
                int nErr = WSAGetLastError();
                if (nErr == WSAEMSGSIZE)
                    continue;
                if (nErr == WSAEWOULDBLOCK || nErr == WSAEINTR || nErr == WSAEINPROGRESS) {
                    this_thread::sleep_for(boost::chrono::milliseconds(10));
                    continue;
                }
            }
            if (!strLine.empty())
                return true;
            if (nBytes == 0) {
                // socket closed
                LogPrintf("IRC: socket closed\n");
                return false;
            } else {
                // socket error
                int nErr = WSAGetLastError();
                LogPrintf("IRC: recv failed: %d\n", nErr);
                return false;
            }
        }
    }
}

string EncodeAddress(const CService& addr) {
    struct ircaddr tmp;
    if (addr.GetInAddr(&tmp.ip))
    {
        tmp.port = htons(addr.GetPort());

        vector<unsigned char> vch(UBEGIN(tmp), UEND(tmp));
        return string("u") + EncodeBase58Check(vch);
    }
    return "";
}

bool DecodeAddress(string str, CService& addr) {
    vector<unsigned char> vch;
    if (!DecodeBase58Check(str.substr(1), vch))
        return false;

    struct ircaddr tmp;
    if (vch.size() != sizeof(tmp))
        return false;
    memcpy(&tmp, &vch[0], sizeof(tmp));

    addr = CService(tmp.ip, ntohs(tmp.port));
    return true;
}

static bool Send(SOCKET hSocket, const char* pszSend) {
    if (strstr(pszSend, "PONG") != pszSend) {
        printf("IRC SENDING: %s\n", pszSend);
    }
    const char* psz = pszSend;
    const char* pszEnd = psz + strlen(psz);
    while (psz < pszEnd) {
        int ret = send(hSocket, psz, pszEnd - psz, MSG_NOSIGNAL);
        if (ret < 0)
            return false;
        psz += ret;
    }
    return true;
}

bool RecvLineIRC(SOCKET hSocket, string& strLine) {
    for (;;) {
    	boost::this_thread::interruption_point();
    	if (ShutdownRequested())
            return false;
        bool fRet = RecvIRCLine(hSocket, strLine);
        if (fRet) {
            vector<string> vWords;
            ParseString(strLine, ' ', vWords);
            if (vWords.size() >= 1 && vWords[0] == "PING") {
                strLine[1] = 'O';
                strLine += '\r';
                Send(hSocket, strLine.c_str());
                continue;
            }
        }
        return fRet;
    }
}

int RecvUntil(SOCKET hSocket, const char* psz1, const char* psz2=NULL, const char* psz3=NULL, const char* psz4=NULL) {
    for (;;) {
    	boost::this_thread::interruption_point();
        string strLine;
        strLine.reserve(10000);
        if (!RecvLineIRC(hSocket, strLine))
            return 0;
        LogPrintf("IRC %s\n", strLine.c_str());
        if (psz1 && strLine.find(psz1) != string::npos)
            return 1;
        if (psz2 && strLine.find(psz2) != string::npos)
            return 2;
        if (psz3 && strLine.find(psz3) != string::npos)
            return 3;
        if (psz4 && strLine.find(psz4) != string::npos)
            return 4;
    }
}

bool Wait(int nSeconds) {
    if (ShutdownRequested()) {
        return false;
    }
    LogPrintf("IRC waiting %d seconds to reconnect\n", nSeconds);
    for (int i = 0; i < nSeconds; i++) {
    	boost::this_thread::interruption_point();
        if (ShutdownRequested()) {
            return false;
        }
        this_thread::sleep_for(boost::chrono::seconds(1));
    }
    return true;
}

bool RecvCodeLine(SOCKET hSocket, const char* psz1, string& strRet) {
    strRet.clear();
    for (;;) {
    	boost::this_thread::interruption_point();
        string strLine;
        if (!RecvLineIRC(hSocket, strLine)) {
            return false;
        }

        vector<string> vWords;
        ParseString(strLine, ' ', vWords);
        if (vWords.size() < 2) {
            continue;
        }

        if (vWords[1] == psz1) {
            LogPrintf("IRC %s\n", strLine.c_str());
            strRet = strLine;
            return true;
        }
    }
}

bool GetIPFromIRC(SOCKET hSocket, string strMyName, CNetAddr& ipRet) {
    Send(hSocket, strprintf("USERHOST %s\r", strMyName.c_str()).c_str());

    string strLine;
    if (!RecvCodeLine(hSocket, "302", strLine)) {
        return false;
    }

    vector<string> vWords;
    ParseString(strLine, ' ', vWords);
    if (vWords.size() < 4) {
        return false;
    }

    string str = vWords[3];
    if (str.rfind("@") == string::npos) {
        return false;
    }
    string strHost = str.substr(str.rfind("@")+1);

    // Hybrid IRC used by lfnet always returns IP when you userhost yourself,
    // but in case another IRC is ever used this should work.
    LogPrintf("GetIPFromIRC() got userhost %s\n", strHost.c_str());
    CNetAddr addr(strHost, true);
    if (!addr.IsValid()) {
        return false;
    }
    ipRet = addr;

    return true;
}

void ThreadIRCSeed() {
    // Make this thread recognisable as the IRC seeding thread
    RenameThread("Ember-ircseed");

    int nErrorWait = 10;
    int nRetryWait = 10;
    int nNameRetry = 0;

begin_irc:
	boost::this_thread::interruption_point();
    CService addrConnect("62.210.131.147", 7777); // irc.lfnet.org

    CService addrIRC("irc.lfnet.org", 7777, true);
    if (addrIRC.IsValid()) {
        addrConnect = addrIRC;
    }

    SOCKET hSocket;
    if (!ConnectSocket(addrConnect, hSocket, 60*5)) {
        LogPrintf("IRC connect failed\n");
        nErrorWait = nErrorWait * 11 / 10;
        if (Wait(nErrorWait += 60))
            goto begin_irc;
        else
            return;
    }

    if (!RecvUntil(hSocket, "Found your hostname", "using your IP address instead", "Couldn't look up your hostname", "ignoring hostname")) {
        closesocket(hSocket);
        hSocket = INVALID_SOCKET;
        nErrorWait = nErrorWait * 11 / 10;
        if (Wait(nErrorWait += 60))
            goto begin_irc;
        else
            return;
    }

    CNetAddr addrIPv4("1.2.3.4"); // arbitrary IPv4 address to make GetLocal prefer IPv4 addresses
    CService addrLocal;
    string strMyName;
    // Don't use our IP as our nick if we're not listening
    // or if it keeps failing because the nick is already in use.
    if (!fNoListen && GetLocal(addrLocal, &addrIPv4) && nNameRetry<3)
        strMyName = EncodeAddress(GetLocalAddress(&addrConnect));
    if (strMyName == "")
        strMyName = strprintf("x%llu", GetRand(1000000000));

    Send(hSocket, strprintf("NICK %s\r", strMyName.c_str()).c_str());
    Send(hSocket, strprintf("USER %s 8 * : %s\r", strMyName.c_str(), strMyName.c_str()).c_str());

    int nRet = RecvUntil(hSocket, " 004 ", " 433 ");
    if (nRet != 1)
    {
        closesocket(hSocket);
        hSocket = INVALID_SOCKET;
        if (nRet == 2) {
            LogPrintf("IRC name already in use\n");
            nNameRetry++;
            Wait(10);
            goto begin_irc;
        }
        nErrorWait = nErrorWait * 11 / 10;
        if (Wait(nErrorWait += 60))
            goto begin_irc;
        else
            return;
    }
    nNameRetry = 0;

    boost::this_thread::interruption_point();
    this_thread::sleep_for(boost::chrono::milliseconds(500));
    boost::this_thread::interruption_point();

    // Get our external IP from the IRC server and re-nick before joining the channel
    CNetAddr addrFromIRC;
    if (GetIPFromIRC(hSocket, strMyName, addrFromIRC)) {
        LogPrintf("GetIPFromIRC() returned %s\n", addrFromIRC.ToString().c_str());
        // Don't use our IP as our nick if we're not listening
        if (!fNoListen && addrFromIRC.IsRoutable()) {
            // IRC lets you to re-nick
            AddLocal(addrFromIRC, LOCAL_IRC);
            strMyName = EncodeAddress(GetLocalAddress(&addrConnect));
            Send(hSocket, strprintf("NICK %s\r", strMyName.c_str()).c_str());
        }
    }

    if (TestNet()) {
        Send(hSocket, "JOIN #Ember0T\r");
        Send(hSocket, "WHO #Ember0T\r");
    } else {
        // randomly join #Ember00-#Ember05
        // int channel_number = GetRandInt(5);

        // Channel number is always 0 for initial release
        int channel_number = 0;
        Send(hSocket, strprintf("JOIN #Ember%02d\r", channel_number).c_str());
        Send(hSocket, strprintf("WHO #Ember%02d\r", channel_number).c_str());
    }

    int64 nStart = GetTime();
    string strLine;
    strLine.reserve(10000);

    while (RecvLineIRC(hSocket, strLine)) {

    	boost::this_thread::interruption_point();

        if (strLine.empty() || strLine.size() > 900 || strLine[0] != ':')
            continue;

        vector<string> vWords;
        ParseString(strLine, ' ', vWords);
        if (vWords.size() < 2)
            continue;

        char pszName[10000];
        pszName[0] = '\0';

        if (vWords[1] == "352" && vWords.size() >= 8) {
            // index 7 is limited to 16 characters
            // could get full length name at index 10, but would be different from join messages
            strlcpy(pszName, vWords[7].c_str(), sizeof(pszName));
            //LogPrintf("IRC got who\n");
        }

        if (vWords[1] == "JOIN" && vWords[0].size() > 1) {
            // :username!username@50000007.F000000B.90000002.IP JOIN :#channelname
            strlcpy(pszName, vWords[0].c_str() + 1, sizeof(pszName));
            if (strchr(pszName, '!'))
                *strchr(pszName, '!') = '\0';
            //LogPrintf("IRC got join\n");
        }

        if (pszName[0] == 'u') {
            CAddress addr;
            if (DecodeAddress(pszName, addr)) {
                addr.nTime = GetAdjustedTime();
                if (addrman.Add(addr, addrConnect, 51 * 60))
                    LogPrintf("IRC got new address: %s\n", addr.ToString().c_str());
                nGotIRCAddresses++;
            } else {
                LogPrintf("IRC decode failed\n");
            }
        }
    }
    closesocket(hSocket);
    hSocket = INVALID_SOCKET;
}

#ifdef TEST
int main(int argc, char *argv[])
{
    WSADATA wsadata;
    if (WSAStartup(MAKEWORD(2,2), &wsadata) != NO_ERROR) {
        LogPrintf("Error at WSAStartup()\n");
        return false;
    }

    ThreadIRCSeed(NULL);

    WSACleanup();
    return 0;
}
#endif
<|endoftext|>


<commit_before>//
// Copyright 2011-2015 Jeff Bush
//
// Licensed 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.
//

//
// The basic approach is based on this article:
// http://www.drdobbs.com/parallel/rasterization-on-larrabee/217200602
// And is also described in "Hierarchical polygon tiling with coverage
// masks" Proceedings of ACM SIGGRAPH 93, Ned Greene.
//

#include "Rasterizer.h"
#include "SIMDMath.h"

using namespace librender;

namespace
{

const veci16_t kXStep = { 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3 };
const veci16_t kYStep = { 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3 };

void setupEdge(int tileLeft, int tileTop, int x1, int y1,
	int x2, int y2, int &outAcceptEdgeValue, int &outRejectEdgeValue,
	veci16_t &outAcceptStepMatrix, veci16_t &outRejectStepMatrix)
{
	veci16_t xAcceptStepValues = kXStep * splati(kTileSize / 4);
	veci16_t yAcceptStepValues = kYStep * splati(kTileSize / 4);
	veci16_t xRejectStepValues = xAcceptStepValues;
	veci16_t yRejectStepValues = yAcceptStepValues;
	int trivialAcceptX = tileLeft;
	int trivialAcceptY = tileTop;
	int trivialRejectX = tileLeft;
	int trivialRejectY = tileTop;
	const int kThreeQuarterTile = kTileSize * 3 / 4;

	if (y2 > y1)
	{
		trivialAcceptX += kTileSize - 1;
		xAcceptStepValues = xAcceptStepValues - splati(kThreeQuarterTile);
	}
	else
	{
		trivialRejectX += kTileSize - 1;
		xRejectStepValues = xRejectStepValues - splati(kThreeQuarterTile);
	}

	if (x2 > x1)
	{
		trivialRejectY += kTileSize - 1;
		yRejectStepValues = yRejectStepValues - splati(kThreeQuarterTile);
	}
	else
	{
		trivialAcceptY += kTileSize - 1;
		yAcceptStepValues = yAcceptStepValues - splati(kThreeQuarterTile);
	}

	int xStep = y2 - y1;
	int yStep = x2 - x1;

	outAcceptEdgeValue = (trivialAcceptX - x1) * xStep - (trivialAcceptY - y1) * yStep;
	outRejectEdgeValue = (trivialRejectX - x1) * xStep - (trivialRejectY - y1) * yStep;

	if (y1 > y2 || (y1 == y2 && x2 > x1))
	{
		// This is a top or left edge.  We adjust the edge equation values by one
		// so it doesn't overlap (top left fill convention).
		outAcceptEdgeValue++;
		outRejectEdgeValue++;
	}

	// Set up xStepValues
	xAcceptStepValues *= splati(xStep);
	xRejectStepValues *= splati(xStep);

	// Set up yStepValues
	yAcceptStepValues *= splati(yStep);
	yRejectStepValues *= splati(yStep);

	// Add together
	outAcceptStepMatrix = xAcceptStepValues - yAcceptStepValues;
	outRejectStepMatrix = xRejectStepValues - yRejectStepValues;
}

// Workhorse of rasterization.  Recursively subdivides tile into 4x4 grids.
void subdivideTile(
	ShaderFiller &filler,
	const int acceptCornerValue1,
	const int acceptCornerValue2,
	const int acceptCornerValue3,
	const int rejectCornerValue1,
	const int rejectCornerValue2,
	const int rejectCornerValue3,
	const veci16_t acceptStep1,
	const veci16_t acceptStep2,
	const veci16_t acceptStep3,
	const veci16_t rejectStep1,
	const veci16_t rejectStep2,
	const veci16_t rejectStep3,
	const int tileSize,
	const int tileLeft,
	const int tileTop,
	const int clipRight,
	const int clipBottom)
{
	// Compute accept masks
	const veci16_t acceptEdgeValue1 = acceptStep1 + splati(acceptCornerValue1);
	const veci16_t acceptEdgeValue2 = acceptStep2 + splati(acceptCornerValue2);
	const veci16_t acceptEdgeValue3 = acceptStep3 + splati(acceptCornerValue3);
	const int trivialAcceptMask = __builtin_nyuzi_mask_cmpi_sle(acceptEdgeValue1, splati(0))
		& __builtin_nyuzi_mask_cmpi_sle(acceptEdgeValue2, splati(0))
		& __builtin_nyuzi_mask_cmpi_sle(acceptEdgeValue3, splati(0));

	if (tileSize == 4)
	{
		// End recursion
		filler.fillMasked(tileLeft, tileTop, trivialAcceptMask);
		return;
	}

	const int subTileSize = tileSize / 4;

	// Process all trivially accepted blocks
	if (trivialAcceptMask != 0)
	{
		int currentMask = trivialAcceptMask;

		while (currentMask)
		{
			const int index = __builtin_clz(currentMask) - 16;
			currentMask &= ~(0x8000 >> index);
			const int subTileLeft = tileLeft + subTileSize * (index & 3);
			const int subTileTop = tileTop + subTileSize * (index >> 2);
			const int right = min(subTileSize, clipRight - subTileLeft);
			const int bottom = min(subTileSize, clipBottom - subTileTop);
			for (int y = 0; y < bottom; y += 4)
			{
				for (int x = 0; x < right; x += 4)
					filler.fillMasked(subTileLeft + x, subTileTop + y, 0xffff);
			}
		}
	}

	// Compute reject masks
	const veci16_t rejectEdgeValue1 = rejectStep1 + splati(rejectCornerValue1);
	const veci16_t rejectEdgeValue2 = rejectStep2 + splati(rejectCornerValue2);
	const veci16_t rejectEdgeValue3 = rejectStep3 + splati(rejectCornerValue3);
	const int trivialRejectMask = __builtin_nyuzi_mask_cmpi_sgt(rejectEdgeValue1, splati(0))
		| __builtin_nyuzi_mask_cmpi_sgt(rejectEdgeValue2, splati(0))
		| __builtin_nyuzi_mask_cmpi_sgt(rejectEdgeValue3, splati(0));

	// Recurse into blocks that are neither trivially rejected or accepted.
	// They are partially overlapped and need to be further subdivided.
	int recurseMask = (trivialAcceptMask | trivialRejectMask) ^ 0xffff;
	if (recurseMask)
	{
		// Divide each step matrix by 4
		const veci16_t subAcceptStep1 = acceptStep1 >> splati(2);
		const veci16_t subAcceptStep2 = acceptStep2 >> splati(2);
		const veci16_t subAcceptStep3 = acceptStep3 >> splati(2);
		const veci16_t subRejectStep1 = rejectStep1 >> splati(2);
		const veci16_t subRejectStep2 = rejectStep2 >> splati(2);
		const veci16_t subRejectStep3 = rejectStep3 >> splati(2);

		while (recurseMask)
		{
			const int index = __builtin_clz(recurseMask) - 16;
			recurseMask &= ~(0x8000 >> index);
			const int x = tileLeft + subTileSize * (index & 3);
			const int y = tileTop + subTileSize * (index >> 2);
			if (x >= clipRight || y >= clipBottom)
				continue;	// Clip tiles that are outside viewport

			subdivideTile(
				filler,
				acceptEdgeValue1[index],
				acceptEdgeValue2[index],
				acceptEdgeValue3[index],
				rejectEdgeValue1[index],
				rejectEdgeValue2[index],
				rejectEdgeValue3[index],
				subAcceptStep1,
				subAcceptStep2,
				subAcceptStep3,
				subRejectStep1,
				subRejectStep2,
				subRejectStep3,
				subTileSize,
				x,
				y,
				clipRight,
				clipBottom);
		}
	}
}

}

void librender::fillTriangle(ShaderFiller &filler,
	int tileLeft, int tileTop,
	int x1, int y1, int x2, int y2, int x3, int y3,
	int clipRight, int clipBottom)
{
	int acceptValue1;
	int rejectValue1;
	veci16_t acceptStepMatrix1;
	veci16_t rejectStepMatrix1;
	int acceptValue2;
	int rejectValue2;
	veci16_t acceptStepMatrix2;
	veci16_t rejectStepMatrix2;
	int acceptValue3;
	int rejectValue3;
	veci16_t acceptStepMatrix3;
	veci16_t rejectStepMatrix3;

	// This assumes counter-clockwise winding for triangles that are
	// facing the camera.
	setupEdge(tileLeft, tileTop, x1, y1, x3, y3, acceptValue1, rejectValue1,
		acceptStepMatrix1, rejectStepMatrix1);
	setupEdge(tileLeft, tileTop, x3, y3, x2, y2, acceptValue2, rejectValue2,
		acceptStepMatrix2, rejectStepMatrix2);
	setupEdge(tileLeft, tileTop, x2, y2, x1, y1, acceptValue3, rejectValue3,
		acceptStepMatrix3, rejectStepMatrix3);

	subdivideTile(
		filler,
		acceptValue1,
		acceptValue2,
		acceptValue3,
		rejectValue1,
		rejectValue2,
		rejectValue3,
		acceptStepMatrix1,
		acceptStepMatrix2,
		acceptStepMatrix3,
		rejectStepMatrix1,
		rejectStepMatrix2,
		rejectStepMatrix3,
		kTileSize,
		tileLeft,
		tileTop,
		clipRight,
		clipBottom);
}
<commit_msg>librender: In some cases, we recurse into a leaf function during rasterization, but there's nothing to do<commit_after>//
// Copyright 2011-2015 Jeff Bush
//
// Licensed 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.
//

//
// The basic approach is based on this article:
// http://www.drdobbs.com/parallel/rasterization-on-larrabee/217200602
// And is also described in "Hierarchical polygon tiling with coverage
// masks" Proceedings of ACM SIGGRAPH 93, Ned Greene.
//

#include "Rasterizer.h"
#include "SIMDMath.h"

using namespace librender;

namespace
{

const veci16_t kXStep = { 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3 };
const veci16_t kYStep = { 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3 };

void setupEdge(int tileLeft, int tileTop, int x1, int y1,
	int x2, int y2, int &outAcceptEdgeValue, int &outRejectEdgeValue,
	veci16_t &outAcceptStepMatrix, veci16_t &outRejectStepMatrix)
{
	veci16_t xAcceptStepValues = kXStep * splati(kTileSize / 4);
	veci16_t yAcceptStepValues = kYStep * splati(kTileSize / 4);
	veci16_t xRejectStepValues = xAcceptStepValues;
	veci16_t yRejectStepValues = yAcceptStepValues;
	int trivialAcceptX = tileLeft;
	int trivialAcceptY = tileTop;
	int trivialRejectX = tileLeft;
	int trivialRejectY = tileTop;
	const int kThreeQuarterTile = kTileSize * 3 / 4;

	if (y2 > y1)
	{
		trivialAcceptX += kTileSize - 1;
		xAcceptStepValues = xAcceptStepValues - splati(kThreeQuarterTile);
	}
	else
	{
		trivialRejectX += kTileSize - 1;
		xRejectStepValues = xRejectStepValues - splati(kThreeQuarterTile);
	}

	if (x2 > x1)
	{
		trivialRejectY += kTileSize - 1;
		yRejectStepValues = yRejectStepValues - splati(kThreeQuarterTile);
	}
	else
	{
		trivialAcceptY += kTileSize - 1;
		yAcceptStepValues = yAcceptStepValues - splati(kThreeQuarterTile);
	}

	int xStep = y2 - y1;
	int yStep = x2 - x1;

	outAcceptEdgeValue = (trivialAcceptX - x1) * xStep - (trivialAcceptY - y1) * yStep;
	outRejectEdgeValue = (trivialRejectX - x1) * xStep - (trivialRejectY - y1) * yStep;

	if (y1 > y2 || (y1 == y2 && x2 > x1))
	{
		// This is a top or left edge.  We adjust the edge equation values by one
		// so it doesn't overlap (top left fill convention).
		outAcceptEdgeValue++;
		outRejectEdgeValue++;
	}

	// Set up xStepValues
	xAcceptStepValues *= splati(xStep);
	xRejectStepValues *= splati(xStep);

	// Set up yStepValues
	yAcceptStepValues *= splati(yStep);
	yRejectStepValues *= splati(yStep);

	// Add together
	outAcceptStepMatrix = xAcceptStepValues - yAcceptStepValues;
	outRejectStepMatrix = xRejectStepValues - yRejectStepValues;
}

// Workhorse of rasterization.  Recursively subdivides tile into 4x4 grids.
void subdivideTile(
	ShaderFiller &filler,
	const int acceptCornerValue1,
	const int acceptCornerValue2,
	const int acceptCornerValue3,
	const int rejectCornerValue1,
	const int rejectCornerValue2,
	const int rejectCornerValue3,
	const veci16_t acceptStep1,
	const veci16_t acceptStep2,
	const veci16_t acceptStep3,
	const veci16_t rejectStep1,
	const veci16_t rejectStep2,
	const veci16_t rejectStep3,
	const int tileSize,
	const int tileLeft,
	const int tileTop,
	const int clipRight,
	const int clipBottom)
{
	// Compute accept masks
	const veci16_t acceptEdgeValue1 = acceptStep1 + splati(acceptCornerValue1);
	const veci16_t acceptEdgeValue2 = acceptStep2 + splati(acceptCornerValue2);
	const veci16_t acceptEdgeValue3 = acceptStep3 + splati(acceptCornerValue3);
	const int trivialAcceptMask = __builtin_nyuzi_mask_cmpi_sle(acceptEdgeValue1, splati(0))
		& __builtin_nyuzi_mask_cmpi_sle(acceptEdgeValue2, splati(0))
		& __builtin_nyuzi_mask_cmpi_sle(acceptEdgeValue3, splati(0));

	if (tileSize == 4)
	{
		// End recursion
		if (trivialAcceptMask)
			filler.fillMasked(tileLeft, tileTop, trivialAcceptMask);

		return;
	}

	const int subTileSize = tileSize / 4;

	// Process all trivially accepted blocks
	if (trivialAcceptMask != 0)
	{
		int currentMask = trivialAcceptMask;

		while (currentMask)
		{
			const int index = __builtin_clz(currentMask) - 16;
			currentMask &= ~(0x8000 >> index);
			const int subTileLeft = tileLeft + subTileSize * (index & 3);
			const int subTileTop = tileTop + subTileSize * (index >> 2);
			const int right = min(subTileSize, clipRight - subTileLeft);
			const int bottom = min(subTileSize, clipBottom - subTileTop);
			for (int y = 0; y < bottom; y += 4)
			{
				for (int x = 0; x < right; x += 4)
					filler.fillMasked(subTileLeft + x, subTileTop + y, 0xffff);
			}
		}
	}

	// Compute reject masks
	const veci16_t rejectEdgeValue1 = rejectStep1 + splati(rejectCornerValue1);
	const veci16_t rejectEdgeValue2 = rejectStep2 + splati(rejectCornerValue2);
	const veci16_t rejectEdgeValue3 = rejectStep3 + splati(rejectCornerValue3);
	const int trivialRejectMask = __builtin_nyuzi_mask_cmpi_sgt(rejectEdgeValue1, splati(0))
		| __builtin_nyuzi_mask_cmpi_sgt(rejectEdgeValue2, splati(0))
		| __builtin_nyuzi_mask_cmpi_sgt(rejectEdgeValue3, splati(0));

	// Recurse into blocks that are neither trivially rejected or accepted.
	// They are partially overlapped and need to be further subdivided.
	int recurseMask = (trivialAcceptMask | trivialRejectMask) ^ 0xffff;
	if (recurseMask)
	{
		// Divide each step matrix by 4
		const veci16_t subAcceptStep1 = acceptStep1 >> splati(2);
		const veci16_t subAcceptStep2 = acceptStep2 >> splati(2);
		const veci16_t subAcceptStep3 = acceptStep3 >> splati(2);
		const veci16_t subRejectStep1 = rejectStep1 >> splati(2);
		const veci16_t subRejectStep2 = rejectStep2 >> splati(2);
		const veci16_t subRejectStep3 = rejectStep3 >> splati(2);

		while (recurseMask)
		{
			const int index = __builtin_clz(recurseMask) - 16;
			recurseMask &= ~(0x8000 >> index);
			const int x = tileLeft + subTileSize * (index & 3);
			const int y = tileTop + subTileSize * (index >> 2);
			if (x >= clipRight || y >= clipBottom)
				continue;	// Clip tiles that are outside viewport

			subdivideTile(
				filler,
				acceptEdgeValue1[index],
				acceptEdgeValue2[index],
				acceptEdgeValue3[index],
				rejectEdgeValue1[index],
				rejectEdgeValue2[index],
				rejectEdgeValue3[index],
				subAcceptStep1,
				subAcceptStep2,
				subAcceptStep3,
				subRejectStep1,
				subRejectStep2,
				subRejectStep3,
				subTileSize,
				x,
				y,
				clipRight,
				clipBottom);
		}
	}
}

}

void librender::fillTriangle(ShaderFiller &filler,
	int tileLeft, int tileTop,
	int x1, int y1, int x2, int y2, int x3, int y3,
	int clipRight, int clipBottom)
{
	int acceptValue1;
	int rejectValue1;
	veci16_t acceptStepMatrix1;
	veci16_t rejectStepMatrix1;
	int acceptValue2;
	int rejectValue2;
	veci16_t acceptStepMatrix2;
	veci16_t rejectStepMatrix2;
	int acceptValue3;
	int rejectValue3;
	veci16_t acceptStepMatrix3;
	veci16_t rejectStepMatrix3;

	// This assumes counter-clockwise winding for triangles that are
	// facing the camera.
	setupEdge(tileLeft, tileTop, x1, y1, x3, y3, acceptValue1, rejectValue1,
		acceptStepMatrix1, rejectStepMatrix1);
	setupEdge(tileLeft, tileTop, x3, y3, x2, y2, acceptValue2, rejectValue2,
		acceptStepMatrix2, rejectStepMatrix2);
	setupEdge(tileLeft, tileTop, x2, y2, x1, y1, acceptValue3, rejectValue3,
		acceptStepMatrix3, rejectStepMatrix3);

	subdivideTile(
		filler,
		acceptValue1,
		acceptValue2,
		acceptValue3,
		rejectValue1,
		rejectValue2,
		rejectValue3,
		acceptStepMatrix1,
		acceptStepMatrix2,
		acceptStepMatrix3,
		rejectStepMatrix1,
		rejectStepMatrix2,
		rejectStepMatrix3,
		kTileSize,
		tileLeft,
		tileTop,
		clipRight,
		clipBottom);
}
<|endoftext|>


<commit_before>/*
 # Copyright (c) 2011, Georgia Tech Research Corporation
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are met:
 #     * Redistributions of source code must retain the above copyright
 #       notice, this list of conditions and the following disclaimer.
 #     * Redistributions in binary form must reproduce the above copyright
 #       notice, this list of conditions and the following disclaimer in the
 #       documentation and/or other materials provided with the distribution.
 #     * Neither the name of the Georgia Tech Research Corporation nor the
 #       names of its contributors may be used to endorse or promote products
 #       derived from this software without specific prior written permission.
 #
 # THIS SOFTWARE IS PROVIDED BY GEORGIA TECH RESEARCH CORPORATION ''AS IS'' AND
 # ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 # DISCLAIMED. IN NO EVENT SHALL GEORGIA TECH BE LIABLE FOR ANY DIRECT, INDIRECT,
 # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
 # OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 # LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 # OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #

 ## author Advait Jain (Healthcare Robotics Lab, Georgia Tech.)
 ## author Chih-Hung Aaron King (Healthcare Robotics Lab, Georgia Tech.)
 */

// This application listens for a rigid body named 'Tracker' on a remote machine.
// The raw data is input to a Extended Kalman Filter (EKF) based estimator estimating
// the target position, velocity, orientation and angular velocity. The estimated
// quantities and raw data are then published through ROS.

#include <iostream>
#include <stdio.h>
#include <math.h>
#include <iostream>

#include <ros/ros.h>
#include <vrpn_Connection.h>
#include <vrpn_Tracker.h>
#include <Eigen/Geometry>
#include <eigen_conversions/eigen_msg.h>
#include <nav_msgs/Odometry.h>
#include <geometry_msgs/TransformStamped.h>
#include <tf/transform_broadcaster.h>
#include <glog/logging.h>

#include "vicon_odometry_estimator.h"

void VRPN_CALLBACK track_target(void *, const vrpn_TRACKERCB tracker);

// NOTE(millanea@ethz.ch): The following callbacks should be implemented in the case that
//                         the tracking system supports them. In this case a flag should be
//                         added to the code indicating the whether or not the vicon estimator
//                         should be run.
//void VRPN_CALLBACK track_target_velocity(void *, const vrpn_TRACKERVELCB tv);
//void VRPN_CALLBACK track_target_acceleration(void *, const vrpn_TRACKERACCCB ta);

// A class representing the state of the tracked target.
class TargetState
{
  public:
    geometry_msgs::TransformStamped measured_transform;
    geometry_msgs::TransformStamped estimated_transform;
    nav_msgs::Odometry estimated_odometry;
};

// Available coordinate systems.
enum CoordinateSystem
{
  vicon,
  optitrack
} coordinate_system;

// Global target descriptions.
TargetState *target_state;
std::string object_name;
std::string coordinate_system_string;

// Global indicating the availability of new VRPN callback function.
bool fresh_data = false;
vrpn_TRACKERCB prev_tracker;

// Pointer to the vicon estimator. Global such that it can be accessed from the callback.
vicon_estimator::ViconOdometryEstimator* vicon_odometry_estimator = NULL;

class Rigid_Body {

  public:
    // Constructor
    Rigid_Body(ros::NodeHandle& nh, std::string server_ip, int port, const std::string& object_name)
    {
      // Advertising published topics.
      measured_target_transform_publisher = nh.advertise<geometry_msgs::TransformStamped>("vicon", 10);
      estimated_target_transform_publisher = nh.advertise<geometry_msgs::TransformStamped>("pose", 10);
      estimated_target_odometry_publisher = nh.advertise<nav_msgs::Odometry>("odometry", 10);
      // Connecting to the vprn device and creating an associated tracker.
      std::stringstream connection_name;
      connection_name << server_ip << ":" << port;
      connection = vrpn_get_connection_by_name(connection_name.str().c_str());
      tracker = new vrpn_Tracker_Remote(object_name.c_str(), connection);
      tracker->print_latest_report();
      // Registering a callback to be called when a new data is made available by the vrpn sever.
      this->tracker->register_change_handler(NULL, track_target);
      tracker->print_latest_report();
      // NOTE(millanea@ethz.ch): The following callbacks should be added if they're available.
      //                         See detailed note above.
      //this->tracker->register_change_handler(NULL, track_target_velocity);
      //this->tracker->register_change_handler(NULL, track_target_acceleration);
    }

    // Publishes the raw measured target state to the transform message.
    void publish_measured_transform(TargetState *target_state)
    {
      measured_target_transform_publisher.publish(target_state->measured_transform);
    }

    // Publishes the estimated target state to the transform message and sends tranform.
    void publish_estimated_transform(TargetState *target_state)
    {
      br.sendTransform(target_state->estimated_transform);
      estimated_target_transform_publisher.publish(target_state->estimated_transform);
    }

    // Publishes the estimated target state to the odometry message.
    void publish_estimated_odometry(TargetState *target_state)
    {
      estimated_target_odometry_publisher.publish(target_state->estimated_odometry);
    }

    // Passes contol to the vrpn client.
    void step_vrpn()
    {
      this->tracker->mainloop();
      this->connection->mainloop();
    }

  private:
    // Publishers
    ros::Publisher measured_target_transform_publisher;
    ros::Publisher estimated_target_transform_publisher;
    ros::Publisher estimated_target_odometry_publisher;
    tf::TransformBroadcaster br;
    // Vprn object pointers
    vrpn_Connection *connection;
    vrpn_Tracker_Remote *tracker;
};

// TODO(millanea@ethz.ch): The following callbacks should be implemented if they're required.
//                         See detailed note above.
//void VRPN_CALLBACK track_target_acceleration(void *, const vrpn_TRACKERACCCB ta) {
//  std::cout<<"acceleration_callback"<<std::endl;
//  std::cout<<"ta.vel[0]"<<ta.acc[0]<<std::endl;
//}
//
//void VRPN_CALLBACK track_target_velocity(void *, const vrpn_TRACKERVELCB tv) {
//  std::cout<<"velocity_callback"<<std::endl;
//  std::cout<<"tv.vel[0]"<<tv.vel[0]<<std::endl;
//}

// Corrects measured target orientation and position for differing frame definitions.
void inline correctForCoordinateSystem(const Eigen::Quaterniond& orientation_in,
                                       const Eigen::Vector3d& position_in,
                                       CoordinateSystem coordinate_system,
                                       Eigen::Quaterniond* orientation_measured_B_W,
                                       Eigen::Vector3d* position_measured_W)
{
  // Rotation to correct between optitrack and vicon
  const Eigen::Quaterniond kqOptitrackFix(0.70710678, 0.70710678, 0.0, 0.0);
  // Correcting measurements based on coordinate system
  switch (coordinate_system)
  {
    case optitrack:
    {
      // Here we rotate the Optitrack measured quaternion by qFix, a
      // Pi/2 rotation around the x-axis. By doing so we convert from
      // NED to ENU.
      *orientation_measured_B_W = kqOptitrackFix * orientation_in * kqOptitrackFix.inverse();
      *position_measured_W = Eigen::Vector3d(position_in.x(), -position_in.z(), position_in.y());
      break;
    }
    case vicon:
    {
      *orientation_measured_B_W = orientation_in;
      *position_measured_W = position_in;
      break;
    }
    default:
    {
      ROS_FATAL("Coordinate system not defined!");
      break;
    }
  }
}

// Compares two instances of tracker data for equality
bool inline tracker_equal(const vrpn_TRACKERCB& vprn_data_1, const vrpn_TRACKERCB& vprn_data_2)
{
  return( vprn_data_1.quat[0] == vprn_data_2.quat[0]
          and vprn_data_1.quat[1] == vprn_data_2.quat[1]
          and vprn_data_1.quat[2] == vprn_data_2.quat[2]
          and vprn_data_1.quat[3] == vprn_data_2.quat[3]
          and vprn_data_1.pos[0] == vprn_data_2.pos[0]
          and vprn_data_1.pos[1] == vprn_data_2.pos[1]
          and vprn_data_1.pos[2] == vprn_data_2.pos[2] );
}

// Tracker Position/Orientation Callback
void VRPN_CALLBACK track_target(void *, const vrpn_TRACKERCB tracker)
{
  // Constructing the raw measured target pose variables.
  Eigen::Quaterniond orientation_in(tracker.quat[3], tracker.quat[0], tracker.quat[1], tracker.quat[2]);
  Eigen::Vector3d position_in(tracker.pos[0], tracker.pos[1], tracker.pos[2]);
  // Constructing the measured pose variables corrected for frame definitions
  Eigen::Quaterniond orientation_measured_B_W;
  Eigen::Vector3d position_measured_W;
  correctForCoordinateSystem(orientation_in, position_in, coordinate_system,
                             &orientation_measured_B_W, &position_measured_W);

  // Verifying that each callback indeed gives fresh data.
  if(tracker_equal(prev_tracker, tracker)) {
    ROS_WARN("Repeated Values");
  }
  prev_tracker = tracker;

  // Somehow the vrpn msgs are in a different time zone.
  const int kMicroSecToNanoSec = 1000;
  ros::Time timestamp_local = ros::Time::now();
  int timediff_sec = std::round(double(timestamp_local.sec - tracker.msg_time.tv_sec) / 3600) * 3600;

  int timestamp_nsec = tracker.msg_time.tv_usec * kMicroSecToNanoSec;
  ros::Time timestamp = ros::Time(tracker.msg_time.tv_sec + timediff_sec, timestamp_nsec);

  ros::Duration time_diff = ros::Time::now() - timestamp;
  if (std::abs(time_diff.toSec()) > 0.1) {
    ROS_WARN_STREAM_THROTTLE(1, "Time delay: " << time_diff.toSec());
  }

  // Updating the estimates with the new measurements.
  vicon_odometry_estimator->updateEstimate(position_measured_W, orientation_measured_B_W);
  vicon_odometry_estimator->publishResults(timestamp);
  Eigen::Vector3d position_estimate_W = vicon_odometry_estimator->getEstimatedPosition();
  Eigen::Vector3d velocity_estimate_W = vicon_odometry_estimator->getEstimatedVelocity();
  Eigen::Quaterniond orientation_estimate_B_W = vicon_odometry_estimator->getEstimatedOrientation();
  Eigen::Vector3d rate_estimate_B = vicon_odometry_estimator->getEstimatedAngularVelocity();

  // Rotating the estimated global frame velocity into the body frame.
  Eigen::Vector3d velocity_estimate_B = orientation_estimate_B_W.toRotationMatrix() * velocity_estimate_W;

  // Populate the raw measured transform message. Published in main loop.
  tf::vectorEigenToMsg(position_measured_W, target_state->measured_transform.transform.translation);
  tf::quaternionEigenToMsg(orientation_measured_B_W, target_state->measured_transform.transform.rotation);

  // Populate the estimated transform message. Published in main loop.
  tf::vectorEigenToMsg(position_estimate_W, target_state->estimated_transform.transform.translation);
  tf::quaternionEigenToMsg(orientation_estimate_B_W, target_state->estimated_transform.transform.rotation);

  // Populate the estimated odometry message. Published in main loop.
  tf::pointEigenToMsg(position_estimate_W, target_state->estimated_odometry.pose.pose.position);
  tf::quaternionEigenToMsg(orientation_estimate_B_W, target_state->estimated_odometry.pose.pose.orientation);
  tf::vectorEigenToMsg(velocity_estimate_B, target_state->estimated_odometry.twist.twist.linear);
  tf::vectorEigenToMsg(rate_estimate_B, target_state->estimated_odometry.twist.twist.angular);

  // Indicating to the main loop the data is ready for publishing.
  fresh_data = true;
}

int main(int argc, char* argv[])
{
  ros::init(argc, argv, "vrpn_client", ros::init_options::AnonymousName);
  ros::NodeHandle nh("");
  ros::NodeHandle private_nh("~");

  target_state = new TargetState;

  std::string vrpn_server_ip;
  int vrpn_port;
  std::string trackedObjectName;

  private_nh.param<std::string>("vrpn_server_ip", vrpn_server_ip, std::string());
  private_nh.param<int>("vrpn_port", vrpn_port, 3883);
  private_nh.param<std::string>("vrpn_coordinate_system", coordinate_system_string, "vicon");
  private_nh.param<std::string>("object_name", object_name, "auk");

  // Debug output
  std::cout << "vrpn_server_ip:" << vrpn_server_ip << std::endl;
  std::cout << "vrpn_port:" << vrpn_port << std::endl;
  std::cout << "vrpn_coordinate_system:" << coordinate_system_string << std::endl;
  std::cout << "object_name:" << object_name << std::endl;

  // Checking for valid coordinate specification
  CHECK((coordinate_system_string == std::string("vicon"))
        || (coordinate_system_string == std::string("optitrack")))
        << "ROS param vrpn_coordinate_system should be either 'vicon' or 'optitrack'!";

  // Creating the estimator
  vicon_odometry_estimator = new vicon_estimator::ViconOdometryEstimator(private_nh);
  vicon_odometry_estimator->initializeParameters(private_nh);
  vicon_odometry_estimator->reset();

  // Creating object which handles data publishing
  Rigid_Body tool(private_nh, vrpn_server_ip, vrpn_port, object_name);

  ros::Rate loop_rate(1000);  //TODO(gohlp): fix this

  while (ros::ok())
  {
    tool.step_vrpn();

    // Publishing newly received data.
    if (fresh_data == true)
    {
      tool.publish_measured_transform(target_state);
      tool.publish_estimated_transform(target_state);
      tool.publish_estimated_odometry(target_state);
      fresh_data = false;
    }
    loop_rate.sleep();
  }
  return 0;
}
<commit_msg>Re-adding removed message time stamping.<commit_after>/*
 # Copyright (c) 2011, Georgia Tech Research Corporation
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are met:
 #     * Redistributions of source code must retain the above copyright
 #       notice, this list of conditions and the following disclaimer.
 #     * Redistributions in binary form must reproduce the above copyright
 #       notice, this list of conditions and the following disclaimer in the
 #       documentation and/or other materials provided with the distribution.
 #     * Neither the name of the Georgia Tech Research Corporation nor the
 #       names of its contributors may be used to endorse or promote products
 #       derived from this software without specific prior written permission.
 #
 # THIS SOFTWARE IS PROVIDED BY GEORGIA TECH RESEARCH CORPORATION ''AS IS'' AND
 # ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 # DISCLAIMED. IN NO EVENT SHALL GEORGIA TECH BE LIABLE FOR ANY DIRECT, INDIRECT,
 # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
 # OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 # LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 # OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #

 ## author Advait Jain (Healthcare Robotics Lab, Georgia Tech.)
 ## author Chih-Hung Aaron King (Healthcare Robotics Lab, Georgia Tech.)
 */

// This application listens for a rigid body named 'Tracker' on a remote machine.
// The raw data is input to a Extended Kalman Filter (EKF) based estimator estimating
// the target position, velocity, orientation and angular velocity. The estimated
// quantities and raw data are then published through ROS.

#include <iostream>
#include <stdio.h>
#include <math.h>
#include <iostream>

#include <ros/ros.h>
#include <vrpn_Connection.h>
#include <vrpn_Tracker.h>
#include <Eigen/Geometry>
#include <eigen_conversions/eigen_msg.h>
#include <nav_msgs/Odometry.h>
#include <geometry_msgs/TransformStamped.h>
#include <tf/transform_broadcaster.h>
#include <glog/logging.h>

#include "vicon_odometry_estimator.h"

void VRPN_CALLBACK track_target(void *, const vrpn_TRACKERCB tracker);

// NOTE(millanea@ethz.ch): The following callbacks should be implemented in the case that
//                         the tracking system supports them. In this case a flag should be
//                         added to the code indicating the whether or not the vicon estimator
//                         should be run.
//void VRPN_CALLBACK track_target_velocity(void *, const vrpn_TRACKERVELCB tv);
//void VRPN_CALLBACK track_target_acceleration(void *, const vrpn_TRACKERACCCB ta);

// A class representing the state of the tracked target.
class TargetState
{
  public:
    geometry_msgs::TransformStamped measured_transform;
    geometry_msgs::TransformStamped estimated_transform;
    nav_msgs::Odometry estimated_odometry;
};

// Available coordinate systems.
enum CoordinateSystem
{
  vicon,
  optitrack
} coordinate_system;

// Global target descriptions.
TargetState *target_state;
std::string object_name;
std::string coordinate_system_string;

// Global indicating the availability of new VRPN callback function.
bool fresh_data = false;
vrpn_TRACKERCB prev_tracker;

// Pointer to the vicon estimator. Global such that it can be accessed from the callback.
vicon_estimator::ViconOdometryEstimator* vicon_odometry_estimator = NULL;

class Rigid_Body {

  public:
    // Constructor
    Rigid_Body(ros::NodeHandle& nh, std::string server_ip, int port, const std::string& object_name)
    {
      // Advertising published topics.
      measured_target_transform_publisher = nh.advertise<geometry_msgs::TransformStamped>("raw_vicon", 10);
      estimated_target_transform_publisher = nh.advertise<geometry_msgs::TransformStamped>("pose", 10);
      estimated_target_odometry_publisher = nh.advertise<nav_msgs::Odometry>("odometry", 10);
      // Connecting to the vprn device and creating an associated tracker.
      std::stringstream connection_name;
      connection_name << server_ip << ":" << port;
      connection = vrpn_get_connection_by_name(connection_name.str().c_str());
      tracker = new vrpn_Tracker_Remote(object_name.c_str(), connection);
      tracker->print_latest_report();
      // Registering a callback to be called when a new data is made available by the vrpn sever.
      this->tracker->register_change_handler(NULL, track_target);
      tracker->print_latest_report();
      // NOTE(millanea@ethz.ch): The following callbacks should be added if they're available.
      //                         See detailed note above.
      //this->tracker->register_change_handler(NULL, track_target_velocity);
      //this->tracker->register_change_handler(NULL, track_target_acceleration);
    }

    // Publishes the raw measured target state to the transform message.
    void publish_measured_transform(TargetState *target_state)
    {
      measured_target_transform_publisher.publish(target_state->measured_transform);
    }

    // Publishes the estimated target state to the transform message and sends tranform.
    void publish_estimated_transform(TargetState *target_state)
    {
      br.sendTransform(target_state->estimated_transform);
      estimated_target_transform_publisher.publish(target_state->estimated_transform);
    }

    // Publishes the estimated target state to the odometry message.
    void publish_estimated_odometry(TargetState *target_state)
    {
      estimated_target_odometry_publisher.publish(target_state->estimated_odometry);
    }

    // Passes contol to the vrpn client.
    void step_vrpn()
    {
      this->tracker->mainloop();
      this->connection->mainloop();
    }

  private:
    // Publishers
    ros::Publisher measured_target_transform_publisher;
    ros::Publisher estimated_target_transform_publisher;
    ros::Publisher estimated_target_odometry_publisher;
    tf::TransformBroadcaster br;
    // Vprn object pointers
    vrpn_Connection *connection;
    vrpn_Tracker_Remote *tracker;
};

// TODO(millanea@ethz.ch): The following callbacks should be implemented if they're required.
//                         See detailed note above.
//void VRPN_CALLBACK track_target_acceleration(void *, const vrpn_TRACKERACCCB ta) {
//  std::cout<<"acceleration_callback"<<std::endl;
//  std::cout<<"ta.vel[0]"<<ta.acc[0]<<std::endl;
//}
//
//void VRPN_CALLBACK track_target_velocity(void *, const vrpn_TRACKERVELCB tv) {
//  std::cout<<"velocity_callback"<<std::endl;
//  std::cout<<"tv.vel[0]"<<tv.vel[0]<<std::endl;
//}

// Corrects measured target orientation and position for differing frame definitions.
void inline correctForCoordinateSystem(const Eigen::Quaterniond& orientation_in,
                                       const Eigen::Vector3d& position_in,
                                       CoordinateSystem coordinate_system,
                                       Eigen::Quaterniond* orientation_measured_B_W,
                                       Eigen::Vector3d* position_measured_W)
{
  // Rotation to correct between optitrack and vicon
  const Eigen::Quaterniond kqOptitrackFix(0.70710678, 0.70710678, 0.0, 0.0);
  // Correcting measurements based on coordinate system
  switch (coordinate_system)
  {
    case optitrack:
    {
      // Here we rotate the Optitrack measured quaternion by qFix, a
      // Pi/2 rotation around the x-axis. By doing so we convert from
      // NED to ENU.
      *orientation_measured_B_W = kqOptitrackFix * orientation_in * kqOptitrackFix.inverse();
      *position_measured_W = Eigen::Vector3d(position_in.x(), -position_in.z(), position_in.y());
      break;
    }
    case vicon:
    {
      *orientation_measured_B_W = orientation_in;
      *position_measured_W = position_in;
      break;
    }
    default:
    {
      ROS_FATAL("Coordinate system not defined!");
      break;
    }
  }
}

// Compares two instances of tracker data for equality
bool inline tracker_equal(const vrpn_TRACKERCB& vprn_data_1, const vrpn_TRACKERCB& vprn_data_2)
{
  return( vprn_data_1.quat[0] == vprn_data_2.quat[0]
          and vprn_data_1.quat[1] == vprn_data_2.quat[1]
          and vprn_data_1.quat[2] == vprn_data_2.quat[2]
          and vprn_data_1.quat[3] == vprn_data_2.quat[3]
          and vprn_data_1.pos[0] == vprn_data_2.pos[0]
          and vprn_data_1.pos[1] == vprn_data_2.pos[1]
          and vprn_data_1.pos[2] == vprn_data_2.pos[2] );
}

// Tracker Position/Orientation Callback
void VRPN_CALLBACK track_target(void *, const vrpn_TRACKERCB tracker)
{
  // Constructing the raw measured target pose variables.
  Eigen::Quaterniond orientation_in(tracker.quat[3], tracker.quat[0], tracker.quat[1], tracker.quat[2]);
  Eigen::Vector3d position_in(tracker.pos[0], tracker.pos[1], tracker.pos[2]);
  // Constructing the measured pose variables corrected for frame definitions
  Eigen::Quaterniond orientation_measured_B_W;
  Eigen::Vector3d position_measured_W;
  correctForCoordinateSystem(orientation_in, position_in, coordinate_system,
                             &orientation_measured_B_W, &position_measured_W);

  // Verifying that each callback indeed gives fresh data.
  if(tracker_equal(prev_tracker, tracker)) {
    ROS_WARN("Repeated Values");
  }
  prev_tracker = tracker;

  // Somehow the vrpn msgs are in a different time zone.
  const int kMicroSecToNanoSec = 1000;
  ros::Time timestamp_local = ros::Time::now();
  int timediff_sec = std::round(double(timestamp_local.sec - tracker.msg_time.tv_sec) / 3600) * 3600;

  int timestamp_nsec = tracker.msg_time.tv_usec * kMicroSecToNanoSec;
  ros::Time timestamp = ros::Time(tracker.msg_time.tv_sec + timediff_sec, timestamp_nsec);

  ros::Duration time_diff = ros::Time::now() - timestamp;
  if (std::abs(time_diff.toSec()) > 0.1) {
    ROS_WARN_STREAM_THROTTLE(1, "Time delay: " << time_diff.toSec());
  }

  // Updating the estimates with the new measurements.
  vicon_odometry_estimator->updateEstimate(position_measured_W, orientation_measured_B_W);
  vicon_odometry_estimator->publishResults(timestamp);
  Eigen::Vector3d position_estimate_W = vicon_odometry_estimator->getEstimatedPosition();
  Eigen::Vector3d velocity_estimate_W = vicon_odometry_estimator->getEstimatedVelocity();
  Eigen::Quaterniond orientation_estimate_B_W = vicon_odometry_estimator->getEstimatedOrientation();
  Eigen::Vector3d rate_estimate_B = vicon_odometry_estimator->getEstimatedAngularVelocity();

  // Rotating the estimated global frame velocity into the body frame.
  Eigen::Vector3d velocity_estimate_B = orientation_estimate_B_W.toRotationMatrix() * velocity_estimate_W;

  // Populate the raw measured transform message. Published in main loop.
  target_state->measured_transform.header.stamp = timestamp;
  target_state->measured_transform.header.frame_id = coordinate_system_string;
  target_state->measured_transform.child_frame_id = object_name;
  tf::vectorEigenToMsg(position_measured_W, target_state->measured_transform.transform.translation);
  tf::quaternionEigenToMsg(orientation_measured_B_W, target_state->measured_transform.transform.rotation);

  // Populate the estimated transform message. Published in main loop.
  target_state->estimated_transform.header.stamp = timestamp;
  target_state->estimated_transform.header.frame_id = coordinate_system_string;
  target_state->estimated_transform.child_frame_id = object_name;
  tf::vectorEigenToMsg(position_estimate_W, target_state->estimated_transform.transform.translation);
  tf::quaternionEigenToMsg(orientation_estimate_B_W, target_state->estimated_transform.transform.rotation);

  // Populate the estimated odometry message. Published in main loop.
  target_state->estimated_odometry.header.stamp = timestamp;
  target_state->estimated_odometry.header.frame_id = coordinate_system_string;
  target_state->estimated_odometry.child_frame_id = object_name;
  tf::pointEigenToMsg(position_estimate_W, target_state->estimated_odometry.pose.pose.position);
  tf::quaternionEigenToMsg(orientation_estimate_B_W, target_state->estimated_odometry.pose.pose.orientation);
  tf::vectorEigenToMsg(velocity_estimate_B, target_state->estimated_odometry.twist.twist.linear);
  tf::vectorEigenToMsg(rate_estimate_B, target_state->estimated_odometry.twist.twist.angular);

  // Indicating to the main loop the data is ready for publishing.
  fresh_data = true;
}

int main(int argc, char* argv[])
{
  ros::init(argc, argv, "vrpn_client", ros::init_options::AnonymousName);
  ros::NodeHandle nh("");
  ros::NodeHandle private_nh("~");

  target_state = new TargetState;

  std::string vrpn_server_ip;
  int vrpn_port;
  std::string trackedObjectName;

  private_nh.param<std::string>("vrpn_server_ip", vrpn_server_ip, std::string());
  private_nh.param<int>("vrpn_port", vrpn_port, 3883);
  private_nh.param<std::string>("vrpn_coordinate_system", coordinate_system_string, "vicon");
  private_nh.param<std::string>("object_name", object_name, "auk");

  // Debug output
  std::cout << "vrpn_server_ip:" << vrpn_server_ip << std::endl;
  std::cout << "vrpn_port:" << vrpn_port << std::endl;
  std::cout << "vrpn_coordinate_system:" << coordinate_system_string << std::endl;
  std::cout << "object_name:" << object_name << std::endl;

  // Checking for valid coordinate specification
  CHECK((coordinate_system_string == std::string("vicon"))
        || (coordinate_system_string == std::string("optitrack")))
        << "ROS param vrpn_coordinate_system should be either 'vicon' or 'optitrack'!";

  // Creating the estimator
  vicon_odometry_estimator = new vicon_estimator::ViconOdometryEstimator(private_nh);
  vicon_odometry_estimator->initializeParameters(private_nh);
  vicon_odometry_estimator->reset();

  // Creating object which handles data publishing
  Rigid_Body tool(private_nh, vrpn_server_ip, vrpn_port, object_name);

  ros::Rate loop_rate(1000);  //TODO(gohlp): fix this

  while (ros::ok())
  {
    tool.step_vrpn();

    // Publishing newly received data.
    if (fresh_data == true)
    {
      tool.publish_measured_transform(target_state);
      tool.publish_estimated_transform(target_state);
      tool.publish_estimated_odometry(target_state);
      fresh_data = false;
    }
    loop_rate.sleep();
  }
  return 0;
}
<|endoftext|>


<commit_before>/*
 * Copyright © 2011-2012 Stéphane Raimbault <stephane.raimbault@gmail.com>
 * Copyright © 2015 Alexandr Kolodkin <alexandr.kolodkin@gmail.com>
 *
 * License ISC, see LICENSE for more details.
 *
 * This library implements the Modbus protocol.
 * http://libmodbus.org/
 *
 */

#include <inttypes.h>

#if defined(ARDUINO) && ARDUINO >= 100
  #include "Arduino.h"
#else
  #include "WProgram.h"
  #include <pins_arduino.h>
#endif
#include "SimpleModbusSlave.h"

#define _MODBUS_RTU_SLAVE                0
#define _MODBUS_RTU_FUNCTION             1
#define _MODBUS_RTU_PRESET_REQ_LENGTH    6
#define _MODBUS_RTU_PRESET_RSP_LENGTH    2

#define _MODBUS_RTU_CHECKSUM_LENGTH      2

#define _MODBUSINO_RTU_MAX_ADU_LENGTH  128

/* Supported function codes */
#define _FC_READ_HOLDING_REGISTERS    0x03
#define _FC_WRITE_MULTIPLE_REGISTERS  0x10

enum {
	_STEP_FUNCTION = 0x01,
	_STEP_META,
	_STEP_DATA
};

static uint16_t crc16(uint8_t *req, uint8_t req_length)
{
	uint8_t j;
	uint16_t crc = 0xFFFF;
	while (req_length--) {
	crc = crc ^ *req++;
	for (j=0; j < 8; j++) {
		if (crc & 0x0001)
			crc = (crc >> 1) ^ 0xA001;
		else
			crc = crc >> 1;
		}
	}
	return (crc << 8  | crc >> 8);
}

SimpleModbusSlave::SimpleModbusSlave(uint8_t slave) {
	if (slave >= 0 & slave <= 247) {
		_slave = slave;
	}
}

void SimpleModbusSlave::setup(long baud) {
	Serial.begin(baud);
}

static int check_integrity(uint8_t *msg, uint8_t msg_length)
{
	uint16_t crc_calculated;
	uint16_t crc_received;

	if (msg_length < 2)
	return -1;

	crc_calculated = crc16(msg, msg_length - 2);
	crc_received = (msg[msg_length - 2] << 8) | msg[msg_length - 1];

	/* Check CRC of msg */
	return (crc_calculated == crc_received) ? msg_length : -1;
}

static int build_response_basis(uint8_t slave, uint8_t function,
				uint8_t* rsp)
{
	rsp[0] = slave;
	rsp[1] = function;

	return _MODBUS_RTU_PRESET_RSP_LENGTH;
}

static void send_msg(uint8_t *msg, uint8_t msg_length)
{
	uint16_t crc = crc16(msg, msg_length);

	msg[msg_length++] = crc >> 8;
	msg[msg_length++] = crc & 0x00FF;

	Serial.write(msg, msg_length);
}

static uint8_t response_exception(uint8_t slave, uint8_t function, uint8_t exception_code, uint8_t *rsp)
{
	uint8_t rsp_length = build_response_basis(slave, function + 0x80, rsp);

	/* Positive exception code */
	rsp[rsp_length++] = exception_code;
	return rsp_length;
}

static void flush(void)
{
	uint8_t i = 0;

	/* Wait a moment to receive the remaining garbage but avoid getting stuck
	 * because the line is saturated */
	while (Serial.available() && i++ < 10) {
		Serial.flush();
		delay(3);
	}
}

static int receive(uint8_t *req, uint8_t _slave)
{
	uint8_t i;
	uint8_t length_to_read;
	uint8_t req_index;
	uint8_t step;
	uint8_t function;

	/* We need to analyse the message step by step.  At the first step, we want
	 * to reach the function code because all packets contain this
	 * information. */
	step = _STEP_FUNCTION;
	length_to_read = _MODBUS_RTU_FUNCTION + 1;

	req_index = 0;
	while (length_to_read != 0) {

	/* The timeout is defined to ~10 ms between each bytes.  Precision is
	   not that important so I rather to avoid millis() to apply the KISS
	   principle (millis overflows after 50 days, etc) */
		if (!Serial.available()) {
	 		i = 0;
				while (!Serial.available()) {
				delay(1);
				if (++i == 10) {
					return -1; /* Too late, bye */
				}
		 	}
		}

		req[req_index] = Serial.read();

		/* Moves the pointer to receive other data */
		req_index++;

		/* Computes remaining bytes */
		length_to_read--;

		if (length_to_read == 0) {
			switch (step) {

			case _STEP_FUNCTION:
				/* Function code position */
				function = req[_MODBUS_RTU_FUNCTION];
				if (function == _FC_READ_HOLDING_REGISTERS) {
					length_to_read = 4;
				} else if (function == _FC_WRITE_MULTIPLE_REGISTERS) {
					length_to_read = 5;
				} else {
					/* Wait a moment to receive the remaining garbage */
					flush();
					if (req[_MODBUS_RTU_SLAVE] == _slave || req[_MODBUS_RTU_SLAVE] == MODBUS_BROADCAST_ADDRESS) {
						/* It's for me so send an exception (reuse req) */
						uint8_t rsp_length = response_exception(_slave, function, MODBUS_EXCEPTION_ILLEGAL_FUNCTION, req);
						send_msg(req, rsp_length);
						return - 1 - MODBUS_EXCEPTION_ILLEGAL_FUNCTION;
					}
					return -1;
				}

				step = _STEP_META;
				break;

			case _STEP_META:
				length_to_read = _MODBUS_RTU_CHECKSUM_LENGTH;

				if (function == _FC_WRITE_MULTIPLE_REGISTERS) {
					length_to_read += req[_MODBUS_RTU_FUNCTION + 5];
				}

				if ((req_index + length_to_read) > _MODBUSINO_RTU_MAX_ADU_LENGTH) {
					flush();
					if (req[_MODBUS_RTU_SLAVE] == _slave || req[_MODBUS_RTU_SLAVE] == MODBUS_BROADCAST_ADDRESS) {
						/* It's for me so send an exception (reuse req) */
						uint8_t rsp_length = response_exception(_slave, function, MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE, req);
						send_msg(req, rsp_length);
						return - 1 - MODBUS_EXCEPTION_ILLEGAL_FUNCTION;
					}
					return -1;
				}
				step = _STEP_DATA;
			}
		}
	}

	return check_integrity(req, req_index);
}

static void reply(uint16_t *tab_reg, uint16_t nb_reg, uint8_t *req, uint8_t req_length, uint8_t _slave)
{
	uint8_t  slave    = req[_MODBUS_RTU_SLAVE];
	uint8_t  function = req[_MODBUS_RTU_FUNCTION];
	uint16_t address  = (req[_MODBUS_RTU_FUNCTION + 1] << 8) + req[_MODBUS_RTU_FUNCTION + 2];
	uint16_t nb       = (req[_MODBUS_RTU_FUNCTION + 3] << 8) + req[_MODBUS_RTU_FUNCTION + 4];
	uint8_t  rsp[_MODBUSINO_RTU_MAX_ADU_LENGTH];
	uint8_t  rsp_length = 0;

	if (slave != _slave && slave != MODBUS_BROADCAST_ADDRESS) {
		return;
	}

	if (address + nb > nb_reg) {
		rsp_length = response_exception(slave, function, MODBUS_EXCEPTION_ILLEGAL_DATA_ADDRESS, rsp);
	} else {
		req_length -= _MODBUS_RTU_CHECKSUM_LENGTH;

		if (function == _FC_READ_HOLDING_REGISTERS) {
			uint16_t i;

			rsp_length = build_response_basis(slave, function, rsp);
			rsp[rsp_length++] = nb << 1;
			for (i = address; i < address + nb; i++) {
				rsp[rsp_length++] = tab_reg[i] >> 8;
				rsp[rsp_length++] = tab_reg[i] & 0xFF;
			}
		} else {
			uint16_t i, j;

			for (i = address, j = 6; i < address + nb; i++, j += 2) {
				/* 6 and 7 = first value */
				tab_reg[i] = (req[_MODBUS_RTU_FUNCTION + j] << 8) +
					req[_MODBUS_RTU_FUNCTION + j + 1];
			}

			rsp_length = build_response_basis(slave, function, rsp);
			/* 4 to copy the address (2) and the no. of registers */
			memcpy(rsp + rsp_length, req + rsp_length, 4);
			rsp_length += 4;
		}
	}

	send_msg(rsp, rsp_length);
}

int SimpleModbusSlave::loop(uint16_t* tab_reg, uint16_t nb_reg)
{
	int rc = 0;
	uint8_t req[_MODBUSINO_RTU_MAX_ADU_LENGTH];

	if (Serial.available()) {
		rc = receive(req, _slave);
		if (rc > 0) {
			reply(tab_reg, nb_reg, req, rc, _slave);
		}
	}

	/* Returns a positive value if successful,
	   0 if a slave filtering has occured,
	   -1 if an undefined error has occured,
	   -2 for MODBUS_EXCEPTION_ILLEGAL_FUNCTION
	   etc */
	return rc;
}
<commit_msg>* Fixed check_integrity function.<commit_after>/*
 * Copyright © 2011-2012 Stéphane Raimbault <stephane.raimbault@gmail.com>
 * Copyright © 2015 Alexandr Kolodkin <alexandr.kolodkin@gmail.com>
 *
 * License ISC, see LICENSE for more details.
 *
 * This library implements the Modbus protocol.
 * http://libmodbus.org/
 *
 */

#include <inttypes.h>

#if defined(ARDUINO) && ARDUINO >= 100
  #include "Arduino.h"
#else
  #include "WProgram.h"
  #include <pins_arduino.h>
#endif
#include "SimpleModbusSlave.h"

#define _MODBUS_RTU_SLAVE                0
#define _MODBUS_RTU_FUNCTION             1
#define _MODBUS_RTU_PRESET_REQ_LENGTH    6
#define _MODBUS_RTU_PRESET_RSP_LENGTH    2

#define _MODBUS_RTU_CHECKSUM_LENGTH      2

#define _MODBUSINO_RTU_MAX_ADU_LENGTH  128

/* Supported function codes */
#define _FC_READ_HOLDING_REGISTERS    0x03
#define _FC_WRITE_MULTIPLE_REGISTERS  0x10

enum {
	_STEP_FUNCTION = 0x01,
	_STEP_META,
	_STEP_DATA
};

static uint16_t crc16(uint8_t *req, uint8_t req_length)
{
	uint8_t j;
	uint16_t crc = 0xFFFF;
	while (req_length--) {
	crc = crc ^ *req++;
	for (j=0; j < 8; j++) {
		if (crc & 0x0001)
			crc = (crc >> 1) ^ 0xA001;
		else
			crc = crc >> 1;
		}
	}
	return (crc << 8  | crc >> 8);
}

SimpleModbusSlave::SimpleModbusSlave(uint8_t slave) {
	if (slave >= 0 & slave <= 247) {
		_slave = slave;
	}
}

void SimpleModbusSlave::setup(long baud) {
	Serial.begin(baud);
}

// Check CRC of msg
static int check_integrity(uint8_t *msg, uint8_t msg_length)
{
	if ((msg_length >= 2) && crc16(msg, msg_length)) {
		return msg_length;
	} else {
		return -1;
	}
}

static int build_response_basis(uint8_t slave, uint8_t function,
				uint8_t* rsp)
{
	rsp[0] = slave;
	rsp[1] = function;

	return _MODBUS_RTU_PRESET_RSP_LENGTH;
}

static void send_msg(uint8_t *msg, uint8_t msg_length)
{
	uint16_t crc = crc16(msg, msg_length);

	msg[msg_length++] = crc >> 8;
	msg[msg_length++] = crc & 0x00FF;

	Serial.write(msg, msg_length);
}

static uint8_t response_exception(uint8_t slave, uint8_t function, uint8_t exception_code, uint8_t *rsp)
{
	uint8_t rsp_length = build_response_basis(slave, function + 0x80, rsp);

	/* Positive exception code */
	rsp[rsp_length++] = exception_code;
	return rsp_length;
}

static void flush(void)
{
	uint8_t i = 0;

	/* Wait a moment to receive the remaining garbage but avoid getting stuck
	 * because the line is saturated */
	while (Serial.available() && i++ < 10) {
		Serial.flush();
		delay(3);
	}
}

static int receive(uint8_t *req, uint8_t _slave)
{
	uint8_t i;
	uint8_t length_to_read;
	uint8_t req_index;
	uint8_t step;
	uint8_t function;

	/* We need to analyse the message step by step.  At the first step, we want
	 * to reach the function code because all packets contain this
	 * information. */
	step = _STEP_FUNCTION;
	length_to_read = _MODBUS_RTU_FUNCTION + 1;

	req_index = 0;
	while (length_to_read != 0) {

	/* The timeout is defined to ~10 ms between each bytes.  Precision is
	   not that important so I rather to avoid millis() to apply the KISS
	   principle (millis overflows after 50 days, etc) */
		if (!Serial.available()) {
	 		i = 0;
				while (!Serial.available()) {
				delay(1);
				if (++i == 10) {
					return -1; /* Too late, bye */
				}
		 	}
		}

		req[req_index] = Serial.read();

		/* Moves the pointer to receive other data */
		req_index++;

		/* Computes remaining bytes */
		length_to_read--;

		if (length_to_read == 0) {
			switch (step) {

			case _STEP_FUNCTION:
				/* Function code position */
				function = req[_MODBUS_RTU_FUNCTION];
				if (function == _FC_READ_HOLDING_REGISTERS) {
					length_to_read = 4;
				} else if (function == _FC_WRITE_MULTIPLE_REGISTERS) {
					length_to_read = 5;
				} else {
					/* Wait a moment to receive the remaining garbage */
					flush();
					if (req[_MODBUS_RTU_SLAVE] == _slave || req[_MODBUS_RTU_SLAVE] == MODBUS_BROADCAST_ADDRESS) {
						/* It's for me so send an exception (reuse req) */
						uint8_t rsp_length = response_exception(_slave, function, MODBUS_EXCEPTION_ILLEGAL_FUNCTION, req);
						send_msg(req, rsp_length);
						return - 1 - MODBUS_EXCEPTION_ILLEGAL_FUNCTION;
					}
					return -1;
				}

				step = _STEP_META;
				break;

			case _STEP_META:
				length_to_read = _MODBUS_RTU_CHECKSUM_LENGTH;

				if (function == _FC_WRITE_MULTIPLE_REGISTERS) {
					length_to_read += req[_MODBUS_RTU_FUNCTION + 5];
				}

				if ((req_index + length_to_read) > _MODBUSINO_RTU_MAX_ADU_LENGTH) {
					flush();
					if (req[_MODBUS_RTU_SLAVE] == _slave || req[_MODBUS_RTU_SLAVE] == MODBUS_BROADCAST_ADDRESS) {
						/* It's for me so send an exception (reuse req) */
						uint8_t rsp_length = response_exception(_slave, function, MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE, req);
						send_msg(req, rsp_length);
						return - 1 - MODBUS_EXCEPTION_ILLEGAL_FUNCTION;
					}
					return -1;
				}
				step = _STEP_DATA;
			}
		}
	}

	return check_integrity(req, req_index);
}

static void reply(uint16_t *tab_reg, uint16_t nb_reg, uint8_t *req, uint8_t req_length, uint8_t _slave)
{
	uint8_t  slave    = req[_MODBUS_RTU_SLAVE];
	uint8_t  function = req[_MODBUS_RTU_FUNCTION];
	uint16_t address  = (req[_MODBUS_RTU_FUNCTION + 1] << 8) + req[_MODBUS_RTU_FUNCTION + 2];
	uint16_t nb       = (req[_MODBUS_RTU_FUNCTION + 3] << 8) + req[_MODBUS_RTU_FUNCTION + 4];
	uint8_t  rsp[_MODBUSINO_RTU_MAX_ADU_LENGTH];
	uint8_t  rsp_length = 0;

	if (slave != _slave && slave != MODBUS_BROADCAST_ADDRESS) {
		return;
	}

	if (address + nb > nb_reg) {
		rsp_length = response_exception(slave, function, MODBUS_EXCEPTION_ILLEGAL_DATA_ADDRESS, rsp);
	} else {
		req_length -= _MODBUS_RTU_CHECKSUM_LENGTH;

		if (function == _FC_READ_HOLDING_REGISTERS) {
			uint16_t i;

			rsp_length = build_response_basis(slave, function, rsp);
			rsp[rsp_length++] = nb << 1;
			for (i = address; i < address + nb; i++) {
				rsp[rsp_length++] = tab_reg[i] >> 8;
				rsp[rsp_length++] = tab_reg[i] & 0xFF;
			}
		} else {
			uint16_t i, j;

			for (i = address, j = 6; i < address + nb; i++, j += 2) {
				/* 6 and 7 = first value */
				tab_reg[i] = (req[_MODBUS_RTU_FUNCTION + j] << 8) +
					req[_MODBUS_RTU_FUNCTION + j + 1];
			}

			rsp_length = build_response_basis(slave, function, rsp);
			/* 4 to copy the address (2) and the no. of registers */
			memcpy(rsp + rsp_length, req + rsp_length, 4);
			rsp_length += 4;
		}
	}

	send_msg(rsp, rsp_length);
}

int SimpleModbusSlave::loop(uint16_t* tab_reg, uint16_t nb_reg)
{
	int rc = 0;
	uint8_t req[_MODBUSINO_RTU_MAX_ADU_LENGTH];

	if (Serial.available()) {
		rc = receive(req, _slave);
		if (rc > 0) {
			reply(tab_reg, nb_reg, req, rc, _slave);
		}
	}

	/* Returns a positive value if successful,
	   0 if a slave filtering has occured,
	   -1 if an undefined error has occured,
	   -2 for MODBUS_EXCEPTION_ILLEGAL_FUNCTION
	   etc */
	return rc;
}
<|endoftext|>